sklearn与pytorch模型的保存与读取

当我们花了很长时间训练了一个模型，需要用该模型做其他事情（比如迁移学习），或者我们想把自己的机器学习模型分享出去的时候，我们这时候需要将我们的ML模型持久化到硬盘中去。

1.sklearn中模型的保存与读取

sklearn模型的保存有两种方法。pickle与joblib。

第一种方法，使用pickle方法保存。

import pickle #pickle模块

#保存模型
with open('model.pickle', 'wb') as f:
    pickle.dump(clf, f)

#读取模型
with open('model.pickle', 'rb') as f:
    model = pickle.load(f)
    #模型读取完毕，之后可以使用该模型做预测
    

第二种方法，使用joblib保存

import joblib

#保存模型
joblib.dump(model, 'model.pkl')

#读取模型
model = joblib.load('model.pkl')
#之后便可以使用该模型做预测了

joblib使用上比较容易，并且读取速度也相对pickle快，因此对于sklearn的模型保存与读取，我强烈推荐使用第二种办法。

2.Pytorch模型的保存与读取

pytorch中模型的保存与读取更更为简单，使用自带的API，load与save即可。

import torch

#保存模型
torch.save(model, "model.pkl")

#读取模型
model = torch.load("model.pkl")
#之后便可以使用该模型进行预测了

太特码简单了 有木有