Packaging Machine Learning Models with Python
Model packaging is an essential step in the Machine Learning deployment process, where the trained model is prepared in a format that can be easily deployed and integrated into production environments. So, if you want to learn about packaging Machine Learning models, this article is for you. In this article, I’ll take you through the task of packaging Machine Learning models with Python.
Packaging Machine Learning Models with Python
Packaging Machine Learning models involves saving it and wrapping it with an API. Here’s the process of packaging a Machine Learning model:
- Step 1: Export/Save the Trained Model
- Step 2: Writing a Wrapper Function to Load the Model
- Step 3: Setting Up an API to Serve the Model
Let’s go through all these steps one by one.
Step 1: Export/Save the Trained Model
Once you have trained a model, the first step is to save it in a format that can be loaded during deployment. Different libraries have different ways of saving models.
Let’s train an example model on the iris dataset and save it using the pickle method:
from sklearn.ensemble import RandomForestClassifier
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
import joblib
# load dataset
iris = load_iris()
X_train, X_test, y_train, y_test = train_test_split(iris.data, iris.target, test_size=0.3)
# train model
model = RandomForestClassifier()
model.fit(X_train, y_train)
# save the model using joblib
joblib.dump(model, 'rf_model.pkl')
Step 2: Writing a Wrapper Function to Load the Model
For deployment, you’ll need a function that loads the model and makes predictions. This wrapper function will be exposed through an API or used in the production environment. Here’s how to write the wrapper function:
# function to load the saved model
def load_model():
model = joblib.load('rf_model.pkl')
return model
# function to make predictions
def predict(input_data):
model = load_model()
prediction = model.predict(input_data)
return prediction
Step 3: Setting Up an API to Serve the Model
To serve your model in production, you can expose it via an API. Flask or FastAPI is often used to create REST APIs for this purpose.
Here’s how to set up an API to serve your model (please make sure to write this piece of code in a separate Python file):
from flask import Flask, request, jsonify
import joblib
import numpy as np
app = Flask(__name__)
# load the model
model = joblib.load('rf_model.pkl')
# define a route for making predictions
@app.route('/predict', methods=['POST'])
def predict():
data = request.get_json() # Get input data from POST request
input_data = np.array(data['input']).reshape(1, -1) # Reshape input
prediction = model.predict(input_data)
return jsonify({'prediction': int(prediction[0])})
if __name__ == '__main__':
app.run(debug=True)
* Serving Flask app '__main__'
* Debug mode: on
INFO:werkzeug:WARNING: This is a development server. Do not use it in a production deployment. Use a production WSGI server instead.
* Running on http://127.0.0.1:5000
INFO:werkzeug:Press CTRL+C to quit
INFO:werkzeug: * Restarting with stat
Now, run this file. Once you run this file, you can test the API using tools like Postman. Postman is a popular API development and testing tool that simplifies the process of sending requests to APIs and inspecting the responses. It is widely used by developers for testing, building, and documenting APIs.
Please download the desktop version of Postman on your system. And follow these steps to test your API:
- Open Postman and create a new POST request.
- Set the URL to http://127.0.0.1:5000/predict (Replace the URL with your URL).
- In the Body tab, choose raw and set the data type to JSON.
Now, enter the following example JSON data in the body (adjust according to your model’s input size):
{
"input": [0.5, 0.3, 0.7, 0.1]
}
// This is an example input for the Machine Learning model.
The image below shows the complete Postman setup and the final output.

Summary
So, this is how we can package Machine Learning models in the form of APIs. Model packaging is an essential step in the machine learning deployment process, where the trained model is prepared in a format that can be easily deployed and integrated into production environments.