【机器学习技巧】机器学习模型的两种存储方式：pickle与joblib模块

本文主要介绍了两种机器学习模型的存储与读取方式，方便我们将训练好的模型直接存储起来，方便下次直接使用该模型进行预测。

1. 构建待存储的示例模型–鸢尾花数据集

# 导入sklearn的官方数据库
from sklearn import datasets
from sklearn.model_selection import train_test_split
from sklearn.model_selection import cross_val_score
# 导入相关的算法
from sklearn.neighbors import KNeighborsClassifier

iris = datasets.load_iris()
# 提取出特征数据
data = iris.data
# 提取标签数据
target = iris.target
# 切分数据：将数据分为训练集和测试集，使用train_test_split方法

# rain_test_split方法首先会把data和target随机打乱，提取出给出比例的数据（包括特征和标签）作为测试数据，下面使用的10%数据作为测试集
x_train, x_test, y_train, y_test = train_test_split(data,target,test_size=0.1)

# 创建模型
knn = KNeighborsClassifier(n_neighbors=5)

# 训练
knn.fit(x_train,y_train)


# 求的值n次交叉验证后的准确率均值
cross_val_score(knn,x_test,y_test,cv=3).mean()  

0.8111111111111112

2. 模型存储

2.1 使用sklearn中的joblib存储

sklearn中提供了高效的模型持久化模块joblib，将模型保存至硬盘

from sklearn.externals import joblib
# 使用joblib.dump保存模型
joblib.dump(knn, 'knn.model')

# 使用joblib.load读取模型
knn_1 = joblib.load('knn.model')

# 使用模型进行预测
y_pre = knn_1.predict(x_test)
print(y_pre)

[1 1 2 1 2 0 1 2 1 2 2 1 2 1 0]

2.2 使用python自带的pickle模块存储模型

import pickle
# pickle.dump保存模型
with open('./knn.pickle', 'wb') as f:
    pickle.dump(knn, f)
    

# pickle.load读取模型
with open('./knn.pickle', 'rb') as f:
    knn2 = pickle.load(f)

    # 使用模型进行预测
y_pre2 = knn2.predict(x_test)
print(y_pre2)

[1 1 2 1 2 0 1 2 1 2 2 1 2 1 0]

如果内容对你有帮助，感谢点赞+关注哦！

欢迎关注我的公众号:阿旭算法与机器学习，共同学习交流。
更多干货内容持续更新中…