Eric Campbell 5 min read Machine Learning

How to Use Pickle for ML Models in Python

Training machine learning models is computationally expensive. Whether you're running a simple logistic regression or a complex ensemble model, you don't want to retrain from scratch every time you...

Key Insights

  • Pickle enables model persistence by serializing trained ML models to disk, eliminating redundant retraining and enabling production deployment scenarios
  • While pickle works seamlessly with scikit-learn models and pipelines, it poses serious security risks when loading untrusted data and suffers from version compatibility issues
  • For production systems, combine pickle with metadata tracking, use joblib for numpy-heavy models, and never deserialize pickle files from untrusted sources

Introduction to Model Persistence

Training machine learning models is computationally expensive. Whether you’re running a simple logistic regression or a complex ensemble model, you don’t want to retrain from scratch every time you need to make predictions. Model persistence solves this problem by saving trained models to disk for later reuse.

Python’s pickle module provides built-in serialization capabilities that work out-of-the-box with most scikit-learn models. It converts Python objects into byte streams that can be stored as files and reconstructed later. This makes pickle the go-to solution for saving ML models during development and for many production scenarios.

The core use cases for model persistence include: avoiding retraining during development iterations, deploying models to production environments, sharing models with team members, and creating reproducible ML pipelines. Pickle handles all these scenarios with minimal code.

Basic Pickle Operations

Pickle operates through two primary functions: pickle.dump() for serialization and pickle.load() for deserialization. Always open files in binary mode ('wb' for writing, 'rb' for reading) since pickle produces binary data.

Here’s a simple example with a Python dictionary:

import pickle

# Create a sample object
data = {
    'model_type': 'classifier',
    'accuracy': 0.95,
    'features': ['age', 'income', 'score']
}

# Save to disk
with open('data.pkl', 'wb') as f:
    pickle.dump(data, f)

# Load from disk
with open('data.pkl', 'rb') as f:
    loaded_data = pickle.load(f)

print(loaded_data)
# Output: {'model_type': 'classifier', 'accuracy': 0.95, 'features': ['age', 'income', 'score']}

The 'wb' and 'rb' modes are crucial. Text mode will corrupt the binary pickle data. Using context managers (with statements) ensures files are properly closed even if errors occur.

Saving and Loading Scikit-learn Models

Scikit-learn models pickle seamlessly because they’re designed with serialization in mind. Here’s a complete workflow with a logistic regression classifier:

import pickle
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score

# Load and split data
iris = load_iris()
X_train, X_test, y_train, y_test = train_test_split(
    iris.data, iris.target, test_size=0.2, random_state=42
)

# Train model
model = LogisticRegression(max_iter=200)
model.fit(X_train, y_train)

# Evaluate before saving
train_accuracy = accuracy_score(y_train, model.predict(X_train))
print(f"Training accuracy: {train_accuracy:.3f}")

# Save the trained model
with open('iris_model.pkl', 'wb') as f:
    pickle.dump(model, f)

# Load the model in a new session
with open('iris_model.pkl', 'rb') as f:
    loaded_model = pickle.load(f)

# Verify it works
test_predictions = loaded_model.predict(X_test)
test_accuracy = accuracy_score(y_test, test_predictions)
print(f"Test accuracy: {test_accuracy:.3f}")

The loaded model retains all learned parameters, hyperparameters, and methods. You can call predict(), predict_proba(), or any other model method exactly as you would with the original object.

Handling Complex ML Pipelines

Real-world ML workflows rarely consist of a single model. You typically have preprocessing steps, feature transformations, and the estimator itself. Scikit-learn’s Pipeline class bundles these components, and pickle handles the entire pipeline as a single object:

import pickle
from sklearn.datasets import load_breast_cancer
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.ensemble import RandomForestClassifier
from sklearn.pipeline import Pipeline

# Load data
data = load_breast_cancer()
X_train, X_test, y_train, y_test = train_test_split(
    data.data, data.target, test_size=0.2, random_state=42
)

# Create pipeline
pipeline = Pipeline([
    ('scaler', StandardScaler()),
    ('classifier', RandomForestClassifier(n_estimators=100, random_state=42))
])

# Train pipeline
pipeline.fit(X_train, y_train)
print(f"Pipeline accuracy: {pipeline.score(X_test, y_test):.3f}")

# Save entire pipeline
with open('cancer_pipeline.pkl', 'wb') as f:
    pickle.dump(pipeline, f)

# Load and use pipeline
with open('cancer_pipeline.pkl', 'rb') as f:
    loaded_pipeline = pickle.load(f)

# The scaler and classifier are both preserved
new_prediction = loaded_pipeline.predict(X_test[:5])
print(f"Predictions: {new_prediction}")

This approach ensures your preprocessing steps are never separated from your model. When you load the pipeline, the scaler’s learned parameters (mean, standard deviation) are restored along with the random forest’s decision trees.

Security Considerations and Limitations

Pickle has a critical security flaw: it can execute arbitrary code during deserialization. A malicious pickle file can run any Python code on your system. Never unpickle data from untrusted sources. This isn’t theoretical—it’s a well-known attack vector.

Version compatibility is another major issue. Pickle files created with one Python or scikit-learn version may fail to load with different versions. This causes problems when deploying models across environments or maintaining models over time.

For large numpy arrays (common in ML models), joblib provides better performance and compression:

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

# Train a model with large arrays
X, y = load_digits(return_X_y=True)
model = RandomForestClassifier(n_estimators=100)
model.fit(X, y)

# joblib is more efficient for numpy-heavy models
joblib.dump(model, 'model.joblib')
loaded_model = joblib.load('model.joblib')

print(f"Accuracy: {loaded_model.score(X, y):.3f}")

Joblib uses optimized serialization for numpy arrays and provides better compression. It’s the recommended approach for scikit-learn models in the official documentation.

Best Practices and Alternatives

Adopt a structured naming convention for model files. Include the model type, date, and version:

import pickle
from datetime import datetime
from sklearn.linear_model import LogisticRegression

model = LogisticRegression()
# ... training code ...

# Descriptive filename
timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
filename = f'logistic_regression_v1_{timestamp}.pkl'

# Save model with metadata
model_data = {
    'model': model,
    'version': '1.0',
    'trained_at': timestamp,
    'accuracy': 0.95,
    'features': ['feature1', 'feature2', 'feature3'],
    'sklearn_version': '1.3.0'
}

with open(filename, 'wb') as f:
    pickle.dump(model_data, f)

# Load and verify metadata
with open(filename, 'rb') as f:
    loaded_data = pickle.load(f)
    print(f"Model version: {loaded_data['version']}")
    print(f"Trained at: {loaded_data['trained_at']}")
    model = loaded_data['model']

For cross-platform deployment or framework interoperability, consider ONNX (Open Neural Network Exchange). For cloud deployments, use platform-specific formats like AWS SageMaker’s model artifacts or Google Cloud’s saved model format.

Practical Deployment Example

Here’s an end-to-end example with a Flask API that serves predictions from a pickled model:

# train_and_save.py
import pickle
from sklearn.datasets import load_iris
from sklearn.ensemble import RandomForestClassifier

# Train model
iris = load_iris()
model = RandomForestClassifier(n_estimators=50, random_state=42)
model.fit(iris.data, iris.target)

# Save for deployment
with open('iris_model.pkl', 'wb') as f:
    pickle.dump(model, f)

print("Model trained and saved successfully")

# api.py
from flask import Flask, request, jsonify
import pickle
import numpy as np

app = Flask(__name__)

# Load model once at startup
with open('iris_model.pkl', 'rb') as f:
    model = pickle.load(f)

@app.route('/predict', methods=['POST'])
def predict():
    # Expect JSON with features
    data = request.get_json()
    features = np.array(data['features']).reshape(1, -1)
    
    prediction = model.predict(features)
    probability = model.predict_proba(features)
    
    return jsonify({
        'prediction': int(prediction[0]),
        'probabilities': probability[0].tolist()
    })

if __name__ == '__main__':
    app.run(debug=True)

This separation of training and serving is standard practice. Train your model in one script, save it, then load it in your production application. The API loads the model once at startup rather than on every request, ensuring low latency.

Pickle remains the simplest solution for model persistence in Python. Use it during development, for internal tools, and in controlled production environments. Just remember its security implications, track your model versions, and consider joblib for numpy-heavy models. For mission-critical production systems, evaluate whether more robust alternatives like ONNX or cloud-native formats better suit your needs.

Liked this? There's more.

Every week: one practical technique, explained simply, with code you can use immediately.