Python-保存已经训练好的机器学习模型

当我们训练好一个model后,下次如果还想用这个model,我们就需要把这个model保存下来,下次直接导入就好了,不然每次都要跑一遍。

1. 使用Python自带的pickle

from sklearn.ensemble import RandomForestClassifier
from sklearn import datasets
import pickle

(X,Y) = datasets.load_iris(return_X_y=True)
rfc = RandomForestClassifier(n_estimators=100,max_depth=100)
rfc.fit(X,Y)
print(rfc.predict(X[0:5,:]))

#保存模型
with open('saved_model/rfc.pickle','wb') as f:
    pickle.dump(rfc,f)
    
#读取模型    
with open('saved_model/rfc.pickle','rb') as f:
    rfc_new = pickle.load(f)
print(rfc_new.predict(X[147:150,:]))

2. 使用sklearn中的joblib模块(速度更快)

from sklearn.ensemble import RandomForestClassifier
from sklearn import datasets
from sklearn.externals import joblib

(X,Y) = datasets.load_iris(return_X_y=True)
rfc = RandomForestClassifier(n_estimators=100,max_depth=100)
rfc.fit(X,Y)
print(rfc.predict(X[0:5,:]))

#保存模型
joblib.dump(rfc,'saved_model/rfc.pkl')

#读取模型
rfc_new = joblib.load('saved_model/rfc.pkl')
print(rfc_new.predict(X[147:150,:]))

 

你可能感兴趣的:(Python-保存已经训练好的机器学习模型)