机器学习 — 训练好的模型保存与加载（joblib、pickle）

机器学习 — 训练好的模型保存与加载（joblib、pickle）

一、前言

当我们训练好一个模型model后，如果在其他程序中或者下次想继续使用这个model，我们就需要把这个model保存下来，下次使用时直接导入就好了，不需要重新训练。

二、实现方式

1. 方式一：joblib

采用joblib模块来保存model，首先安装joblib:

pip install joblib

核心代码如下

import joblib

# 保存 model
joblib.dump(regr, '../../model/regr.pkl')

# 加载 model
clf = joblib.load('../../model/regr.pkl')

完整示例代码如下：

#!/usr/bin/env python
# encoding: utf-8
'''
@Author  : pentiumCM
@Email   : *********@qq.com
@Software: PyCharm
@File    : linear_regression.py
@Time    : 2019/12/19 0:19
@desc	 : 线性回归示例
'''

from matplotlib import pyplot as plot
import numpy as np
from sklearn import linear_model
import joblib as jl

# 构建训练数据
X = np.array([48, 57, 50, 54, 64, 61, 43, 59])              # 体重
Y = np.array([165, 165, 157, 170, 175, 165, 155, 170])      # 身高

# 调整训练数据的维度
X_train = np.reshape(X, (-1, 1))
Y_train = np.reshape(Y, (-1, 1))

# 建立线性模型
regr = linear_model.LinearRegression()
regr.fit(X_train, Y_train)

# 保存 model
jl.dump(regr, '../../model/regr.pkl')

print("模型保存成功")

# 加载 model
clf = jl.load('../../model/regr.pkl')
Y_pred = clf.predict(X_train)

# 可视化
# 1.训练集数据
plot.scatter(X_train, Y_train, color='green')
# 2.线性预测数据
plot.plot(X_train, Y_pred, color='red')
plot.show()

关于pkl文件：
生成的model文件是pkl文件，pkl是python的一种存储文件，需要安装python打开，读取pkl文件需要import pickle包。
所以下面介绍直接使用python自带的pickle来保存训练好的model

2. 方式二：pickle

使用python自带的pickle，来保存和加载模型，核心代码如下

import pickle

# 保存 model
f = open('../../model/regr2.pkl', 'wb')
pickle.dump(regr, f)
f.close()

# 加载 model
f = open('../../model/regr2.pkl', 'rb')
clf = pickle.load(f)
f.close()

完整示例代码如下：

#!/usr/bin/env python
# encoding: utf-8
'''
@Author  : pentiumCM
@Email   : *********@qq.com
@Software: PyCharm
@File    : linear_regression.py
@Time    : 2019/12/19 0:19
@desc	 : 线性回归示例
'''

from matplotlib import pyplot as plot
import numpy as np
from sklearn import linear_model
import pickle

# 构建训练数据
X = np.array([48, 57, 50, 54, 64, 61, 43, 59])              # 体重
Y = np.array([165, 165, 157, 170, 175, 165, 155, 170])      # 身高

# 调整训练数据的维度
X_train = np.reshape(X, (-1, 1))
Y_train = np.reshape(Y, (-1, 1))

# 建立线性模型
regr = linear_model.LinearRegression()
regr.fit(X_train, Y_train)

# 保存 model
f = open('../../model/regr2.pkl', 'wb')
pickle.dump(regr, f)
f.close()

print("模型保存成功")

# 加载 model
f = open('../../model/regr2.pkl', 'rb')
clf = pickle.load(f)
f.close()

Y_pred = clf.predict(X_train)

# 可视化
# 1.训练集数据
plot.scatter(X_train, Y_train, color='green')
# 2.测试数据
plot.plot(X_train, Y_pred, color='red')
plot.show()

注意：
pickle.dump(regr, f)和pickle.load(f)要一起使用，不然会报以下问题：
_pickle.UnpicklingError: invalid load key, ‘\x00’.