如何保存sklearn训练得到的模型？看这一篇就够了

Posted 2022-02-25 serity

目录

• 前言
• 一、安装joblib
• 二、模型的保存与加载
• 参考

前言

很多时候，我们需要用大量的时间来训练一个模型，这些模型如果不加以保存，下次启动电脑时又需要重新进行训练，这无疑是浪费时间的，那么有没有一种方法它能够把训练得到的模型给保存下来呢？

答案是 joblib。

一、安装joblib

有以下两种方式。

pip：

pip install joblib

conda：

conda install joblib

二、模型的保存与加载

joblib 可以保存Python中的任意对象，相关方法如下：

方法	描述
joblib.dump(object, filename)	filename为字符串，且文件后缀名须为 .joblib；用于保存对象
joblib.load(filename)	用于加载对象

例如，我们可以保存一个NumPy的数组：

import numpy as np
import joblib

A = np.arange(4)
joblib.dump(A, 'array.joblib')

然后再加载它：

B = joblib.load('array.joblib')
print(B)
# [0 1 2 3]

再例如，对于训练得到的SVM模型，我们可以将其保存下来以便多次使用：

from sklearn.model_selection import train_test_split
from sklearn.datasets import make_blobs
from sklearn.svm import SVC
import joblib

X, y = make_blobs(n_samples=500, centers=2, random_state=34)
X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=42)
clf = SVC()
clf.fit(X_train, y_train)
joblib.dump(clf, 'svm.joblib')

svm = joblib.load('svm.joblib')
print(svm.score(X_test, y_test))
# 0.92

参考

[1] 官方文档