Skip to content

Save and Load Machine Learning Models

Once you're satisfied with your machine learning model's performance, you can use this template to save your model to a file and, later, load it to make predictions. To do so, this template uses the pickle module, which serializes Python objects. This is useful if you want to send your model to others or if your model has a long training time. When you load your model, you won't need to retrain it to make predictions.

# Loading packages
from sklearn import datasets, model_selection,tree
import pickle

Save your model using pickle

To demonstrate the technique, we'll train a decision tree on the commonly-used iris dataset and save the resulting model as a pickle file.

# Insert the code for your model below
X, y = datasets.load_iris(return_X_y=True)
X_train, X_test, y_train, y_test = model_selection.train_test_split(X, y)
model = tree.DecisionTreeClassifier(max_depth=11, min_samples_leaf=2)
model.fit(X_train, y_train)

# Saving the model as a pickle file
pickle.dump(
    model,  # Replace with your model
    open("model_v1.pickle", "wb"),  # Specify where to save the model
)

If the cell above runs with no errors, you should find the model in the file path you designated.

Load your model using pickle

# Load a pickle file
pickled_model = pickle.load(
    open("model_v1.pickle", "rb")  # Specify which file to load
)

pickled_model

Now you can use pickle_model as you usually use a trained scikit-learn model. For example, you can feed it data and output predictions:

# Output predictions on test set
pickled_model.predict(X_test)