How to Save Trained Model in Python

When working on real-world machine learning (ML) use cases, finding the best algorithm/model is not the end of your responsibilities. It is crucial to save, store, and package these models for their future use and deployment to production.

These practices are needed for a number of reasons:

• Backup: A trained model can be saved as a backup in case the original data is damaged or destroyed. 

• Reusability & reproducibility: Building ML models is time-consuming by nature. To save cost and time, it becomes essential that your model gets you the same results every time you run it. Saving and storing your model the right way takes care of this.

• Deployment: When deploying a trained model in a real-world setting, it becomes necessary to package it for easy deployment. This makes it possible for other systems and applications to use the same model without much hassle.

To reiterate, while saving and storing ML models allow ease of sharing, reusability, and reproducibility; packaging the models enables quick and painless deployment. These 3 operations work in harmony to simplify the whole model management process. 

In this article, you will learn about different methods of saving, storing, and packaging a trained machine-learning model, along with the pros and cons of each method. But before that, you must understand the distinction between these three terms. 

Save vs package vs store ML models

Although all these terms look similar, they are not the same. 

Saving a model refers to the process of saving the model’s parameters, weights, etc., to a file. Usually, all ML and DL models provide some kind of method (eg. model.save()) for saving the models. But you must be aware that save is a single action and gives only a model binary file, so you still need code to make your ML application production-ready.

Packaging, on the other hand, refers to the process of bundling or containerizing the necessary components of a model, such as the model file, dependencies, configuration files, etc., into a single deployable package. The goal of a package is to make it easier to distribute and deploy the ML model in a production environment. 

Once packaged, a model can be deployed across different environments, which allows the model to be used in various production settings such as web applications, mobile applications, etc. Docker is one of the tools which allows you to do this.

Storing the ML model refers to the process of saving the trained model files in a centralized storage that can be accessed anytime when needed. When storing a model, you normally choose some sort of storage from where you can fetch your model and use it anytime. The model registry is a category of tools that solve this issue for you.

Now let’s see how we can save our model.

How to save a trained model in Python?

In this section, you will see different ways of saving machine learning (ML) as well as deep learning (DL) models. To begin with, let’s create a simple classification model using the most famous Iris-dataset. 

Note: The focus of this article is not to show you how you can create the best ML model but to explain how effectively you can save trained models. 

You first need to load the required dependencies and the iris dataset as follows:

# load dependencies
import pandas as pd 

from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler 
from sklearn.neighbors import KNeighborsClassifier
from sklearn.metrics import classification_report, confusion_matrix

# load the dataset
url = "iris.data"

# column names to use
names = ['sepal-length', 'sepal-width', 'petal-length', 'petal-width', 'Class']

# read the dataset from the URL
dataset = pd.read_csv(url, names=names) 

# check the first few rows of iris-classification data
dataset.head()

Next, you need to split the data into training and testing sets and apply the required preprocessing stages, such as feature standardization. 

# separate the independent and dependent features
X = dataset.iloc[:, :-1].values
y = dataset.iloc[:, 4].values 

# Split dataset into random training and testing subsets
X_train, X_test, y_train, y_test = train_test_split(X, 
                                                    y, test_size=0.20) 
# feature standardization
scaler = StandardScaler()
scaler.fit(X_train)

X_train = scaler.transform(X_train)
X_test = scaler.transform(X_test) 

Finally, you need to train a classification model (feel free to choose any) on training data and check its performance on testing data. 

# training a KNN classifier
model = KNeighborsClassifier(n_neighbors=5)
model.fit(X_train, y_train) 

# make predictions on the testing data
y_predict = model.predict(X_test)

# check results
print(confusion_matrix(y_test, y_predict))
print(classification_report(y_test, y_predict)) 

Now you have an ML model that you want to save for future use. The first way to save an ML model is by using the pickle file. 

Saving trained model with pickle

The pickle module can be used to serialize and deserialize the Python objects. Pickling is the process of converting a Python object hierarchy into a byte stream, while Unpickling is the process of converting a byte stream (from a binary file or other object that appears to be made of bytes) back to an object hierarchy.

For saving the ML models used as a pickle file, you need to use the Pickle module that already comes with the default Python installation. 

To save your iris classifier model you simply need to decide on a filename and dump your model to a pickle file like this:

import pickle

# save the iris classification model as a pickle file
model_pkl_file = "iris_classifier_model.pkl"  

with open(model_pkl_file, 'wb') as file:  
    pickle.dump(model, file)

As you can see the file is opened in wb (write binary) mode for saving the model as bytes. Also, the dump() method stores the model in the given pickle file. 

You can also load this model using the load() method of the pickle module. Now you need to open the file in rb (read binary) mode to load the saved model.

# load model from pickle file
with open(model_pkl_file, 'rb') as file:  
    model = pickle.load(file)

# evaluate model 
y_predict = model.predict(X_test)

# check results
print(classification_report(y_test, y_predict)) 

Once loaded you can use this model to make predictions. 

Pros of the Python pickle approach 

Cons of the Python Pickle Approach

Pickle is most suited for small-size models and also has some security issues, these reasons are enough to look for another alternative for saving the ML models. Next, let’s discuss Joblib to save and load ML models. 

Note: In the upcoming sections you will see the same iris classifier model to be saved using different techniques. 

Saving trained model with Joblib

Joblib is a set of tools (typically part of the Scipy ecosystem) that provide lightweight pipelining in Python. It majorly focuses on disk-caching, memoization, and parallel computing and is used for saving and loading Python objects. Joblib has been specifically optimized for NumPy arrays to make it fast and reliable for ML models that have a lot of parameters.

To save large models with Joblib, you need to use the Python Joblib module that comes preinstalled with Python.

import joblib 

# save model with joblib 
filename = 'joblib_model.sav'
joblib.dump(model, filename)

To save the model, you need to define a filename with a ‘.sav’ or ‘.pkl’ extension and call the dump() method from Joblib. 

Similar to pickle, Joblib provides the load() method to load the saved ML model. 

# load model with joblib
loaded_model = joblib.load(filename)

# evaluate model 
y_predict = model.predict(X_test)

# check results
print(classification_report(y_test, y_predict)) 

After loading the model with Joblib you are free to use it on the data to make predictions. 

Pros of saving ML models with Joblib 

Cons of Saving ML Models with Joblib 

Although Joblib solves the major issues faced by pickle, it has some issues on its own. Next, you will see how you can manually save and restore the models using JSON. 

Saving trained model with JSON

When you want to have full control over the save and restore procedure of your ML model, JSON comes into play. Unlike the other two methods, this method does not directly dump the ML model to a file; instead, you need to explicitly define the different parameters of your model to save them. 

To use this method, you need to use the Python json module that again comes along with the default Python installation. Using the JSON method requires additional effort to write all parameters that an ML model contains. To save the model using JSON, let’s create a function like this: 

import json 

# create json save function
def save_json(model, filepath, X_train, y_train): 
    saved_model = {}
    saved_model["algorithm"] = model.get_params()['algorithm'],
    saved_model["max_iter"] = model.get_params()['leaf_size'],
    saved_model["solver"] = model.get_params()['metric'],
    saved_model["metric_params"] = model.get_params()['metric_params'],
    saved_model["n_jobs"] = model.get_params()['n_jobs'],
    saved_model["n_neighbors"] = model.get_params()['n_neighbors'],
    saved_model["p"] = model.get_params()['p'],
    saved_model["weights"] = model.get_params()['weights'],
    saved_model["X_train"] = X_train.tolist() if X_train is not None else "None",
    saved_model["y_train"] = y_train.tolist() if y_train is not None else "None"
    
    json_txt = json.dumps(saved_model, indent=4)
    with open(filepath, "w") as file: 
        file.write(json_txt)

# save the iris-classification model in a json file
file_path = 'json_model.json'
save_json(model, file_path, X_train, y_train)

You see how you need to define each model parameter and the data to store it in JSON. Different models have different methods to check out the parameter details. For example, the get_params() for KNeighboursClassifier gives the list of all the hyperparameters in the model. You need to save all these hyperparameters and data values in a dictionary which is then dumped into a file with the ‘.json’ extension. 

To read this JSON file you just need to open it and access the parameters as follows:

# create json load function 
def load_json(filepath): 
    with open(filepath, "r") as file:
        saved_model = json.load(file)
    
    return saved_model

# load model configurations
saved_model = load_json('json_model.json')
saved_model

In the above code, a function load_json() is created that opens the JSON file in read mode and returns all the parameters and data as a dictionary. 

Unfortunately, you can not use the saved model directly with JSON, you need to read these parameters and data to retrain the model all by yourself. 

Pros of saving ML models with JSON 

Cons of Saving ML Models with JSON

To ensure security and JSON/pickle benefits, you can save your model to a dedicated database. Next, you will see how you can save an ML model in a database. 

Saving deep learning model with TensorFlow Keras

TensorFlow is a popular framework for training DL-based models, and Keras is a wrapper for TensorFlow. A neural network design with numerous layers and a set of labeled data are used to train deep learning models. These models have two major components, Weights and Network architecture, that you need to save to restore them for future use. Typically there are two ways to save deep learning models:

1. Save the model architecture in a JSON or YAML file and weights in an HDF5 file. 
2. Save both model and architecture both in HDF5, protobuf, or tflite file. 

You can refer to any one way to do this, but the widely used method is to save the model weights and architecture together in an HDF5 file. 

To save a deep learning model in TensorFlow Keras, you can use the save() method of the Keras Model object. This method saves the entire model, including the model architecture, optimizer, and weights, in a format that can be loaded later to make predictions.

Here’s an example code snippet that shows how to save a TensorFlow Keras-based DL model:

# import tensorflow dependencies
from tensorflow.keras.models import Sequential, model_from_json
from tensorflow.keras.layers import Dense

# define model architecture
model = Sequential()
model.add(Dense(12, input_dim=4, activation='relu'))
model.add(Dense(8, activation='relu'))
model.add(Dense(1, activation='sigmoid'))

# Compile model
model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['accuracy'])

# Fit the model
model.fit(X_train, y_train, epochs=150, batch_size=10, verbose=0)

# save model and its architecture 
model.save('model.h5')

This is it, you just need to define the model architecture, train the models with appropriate settings, and finally save it using the save() method. 

Loading the saved models with Keras is as easy as reading a file in Python. You just need to call the load_model() method by providing the model file path and your model will be loaded. 

# define dependency 
from tensorflow.keras.models import load_model

# load model 
model = load_model('model.h5')

# check model info 
model.summary()

Your model is now loaded for use. 

Pros of saving models with TensorFlow Keras 

Cons of Saving Models with TensorFlow Keras 

There is one more widely used framework named Pytorch for training the DL-based models. Let’s check how you can save Pytorch-based deep learning models with Python. 

Saving deep learning model with Pytorch

Developed by Facebook, Pytorch is one of the highly used frameworks for developing DL-based solutions. It provides a dynamic computational graph, which allows you to modify your model on-the-fly, making it ideal for research and experimentation. It uses ‘.pt’ and ‘.pth’ file formats to save model architecture and its weights. 

To save a deep learning model in PyTorch, you can use the save() method of the PyTorch torch.nn.Module object. This method saves the entire model, including the model architecture and weights, in a format that can be loaded later to make predictions.

Here’s an example code snippet that shows how to save a PyTorch model:

# import dependencies
import torch
import torch.nn as nn
import numpy as np

# convert data numpy arrays to tensors
X_train = torch.FloatTensor(X_train)
X_test = torch.FloatTensor(X_test)
y_train = torch.LongTensor(y_train)
y_test = torch.LongTensor(y_test)

# define model architecture
class NeuralNetworkClassificationModel(nn.Module):
    def __init__(self,input_dim,output_dim):
        super(NeuralNetworkClassificationModel,self).__init__()
        self.input_layer    = nn.Linear(input_dim,128)
        self.hidden_layer1  = nn.Linear(128,64)
        self.output_layer   = nn.Linear(64,output_dim)
        self.relu = nn.ReLU()
    
    
    def forward(self,x):
        out =  self.relu(self.input_layer(x))
        out =  self.relu(self.hidden_layer1(out))
        out =  self.output_layer(out)
        return out

# define input and output dimensions
input_dim  = 4 
output_dim = 3
model = NeuralNetworkClassificationModel(input_dim,output_dim)

# create our optimizer and loss function object
learning_rate = 0.01
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(),lr=learning_rate)

# define training steps
def train_network(model,optimizer,criterion,X_train,y_train,X_test,y_test,num_epochs,train_losses,test_losses):
    for epoch in range(num_epochs):
        # clear out the gradients from the last step loss.backward()
        optimizer.zero_grad()
        
        # forward feed
        output_train = model(X_train)

        # calculate the loss
        loss_train = criterion(output_train, y_train)

        # backward propagation: calculate gradients
        loss_train.backward()

        # update the weights
        optimizer.step()
        
        output_test = model(X_test)
        loss_test = criterion(output_test,y_test)

        train_losses[epoch] = loss_train.item()
        test_losses[epoch] = loss_test.item()

        if (epoch + 1) % 50 == 0:
            print(f"Epoch { epoch+1 }/{ num_epochs }, Train Loss: { loss_train.item():.4f }, Test Loss: {loss_test.item():.4f}")

# train model
num_epochs = 1000
train_losses = np.zeros(num_epochs)
test_losses  = np.zeros(num_epochs)
train_network(model,optimizer,criterion,X_train,y_train,X_test,y_test,num_epochs,train_losses,test_losses)

# save model 
torch.save(model, 'model_pytorch.pt')

Unlike Tensorflow, Pytorch allows you to have more control over the model training, as seen in the above code. After training the model, you can save the weights and their architecture using save() method. 

Loading the saved model with Pytorch requires the use of load() method. 

# load model
model = torch.load('model_pytorch.pt')
# check model summary
model.eval()

Pros of saving models with Pytorch 

Cons of Saving Models with Pytorch 

This isn’t it, you can use model registry platforms to save DL-based models as well, specially the ones with large size. This makes it easy to deploy and maintain them without requiring extra effort from developers. 

You can find the dataset and code used in this article here. 

How to package ML models?

An ML model is typically optimized for performance on the training dataset and the specific environment in which it is trained. But, when it comes to deploying the models in different environments, such as a production environment, there could be various challenges.

These challenges are but not limited to differences in hardware, software, and data inputs. Packaging the model makes it easier to address these problem, as it allows the model to be exported or serialized into a standard format that can be loaded and used in various environments.

There are various options available for packaging right now. By packaging the model in a standard format such as PMML (Predictive Model Markup Language), ONNX, TensorFlow SavedModel format, etc. it becomes easier to share and collaborate on a model without being concerned about different libraries and tools used by different teams. Now, let’s check a few examples of packaging an ML model with different frameworks in Python.

Note: For this section as well, you will see the same iris-classification example.

Packaging models with PMML

Using the PMML library in Python, you can export your machine learning models to PMML format and then deploy that as a web service, a batch processing system, or a data integration platform. This can make it easier to share and collaborate on machine learning models, as well as to deploy them in various production environments.

To package an ML model using PMML you can use different modules like sklearn2pmml, jpmml-sklearn, jpmml-tensorflow, etc.

Note: To use PMML, you must have Java Runtime installed on your system.

Here is an example code snippet that allows you to package the trained iris classifier model using PMML. 

from sklearn2pmml import PMMLPipeline, sklearn2pmml
# package iris classifier model with PMML
sklearn2pmml(PMMLPipeline([("estimator",
                        	model)]),
         	"iris_model.pmml",
         	with_repr=True)

In the above code, you simply need to create a PMML pipeline object by passing your model object. Then you need to save the PMML object using sklearn2pmml() method. That is it, now you can use this “iris_model.pmml” file across different environments.  

Pros of using PMML 

Cons of using PMML

Packaging models with ONNX

Developed by Microsoft and Facebook, ONNX (Open Neural Network Exchange) is an open format for representing machine learning models. It allows for interoperability between different deep-learning frameworks and tools. 

ONNX models can be deployed efficiently on a variety of platforms, including mobile devices, edge devices, and the cloud. It supports a variety of runtimes, including Caffe2, TensorFlow, PyTorch, and MXNet, which allows you to deploy your models on different devices and platforms with minimal effort.

To save the model using ONNX, you need to have onnx and onnxruntime packages downloaded in your system.

Here is an example of how you can convert the existing ML model to ONNX format.

# load dependencies
import onnxmltools
import onnxruntime

# Convert the KNeighborsClassifier model to ONNX format
onnx_model = onnxmltools.convert_sklearn(model)

# Save the ONNX model in a file
onnx_file = "iris_knn.onnx"
onnxmltools.utils.save_model(onnx_model, onnx_file)

You just need to import the required modules and use the convert_sklearn() method to corvet the sklearn model to the ONNX model. Once the conversion is done, using the save_model() method, you can store the ONNX model in a file with the “.onnx” extension. Although here you see an example of an ML model, ONNX is majorly used for DL models. 

You can also load this model using the ONNX Runtime module.

# Load the ONNX model into ONNX Runtime
sess = onnxruntime.InferenceSession(onnx_file)

# Evaluate the model on some test data
input_data = {"X": X_test[:10].astype('float32')}
output = sess.run(None, input_data)

You need to create a session using InferenceSession() method to load the ONNX model from a file and then use sess.run() method to make predictions from the model. 

Pros of using ONNX

Cons of using ONNX 

Packaging models with Tensorflow SavedModel

Tensorflow’s SavedModel format allows you to easily save and load your deep learning models, and it ensures compatibility with other Tensorflow tools and platforms. Additionally, it provides a streamlined and efficient way to deploy our models in production environments. 

SavedModel supports a wide range of deployment scenarios, including serving models with Tensorflow Serving, deploying models to mobile devices with Tensorflow Lite, and exporting models to other ML libraries such as ONNX.

 It provides a simple and streamlined way to save and load Tensorflow models. The API is easy to use and well-documented, and the format is designed to be efficient and scalable.

Note: You can use the same TensorFlow model trained in the above section.

To save the model in SavedModel format, you can use the following lines of code:

import tensorflow as tf

# using SavedModel format to save the model
tf.saved_model.save(model, "my_model")

You can also load the model with load() method. 

# Load the model
loaded_model = tf.saved_model.load("my_model")

Pros of using Tensorflow SavedModel 

Cons of using Tensorflow SavedModel

Now that you have seen multiple ways of packaging ML and DL models, you must also be aware that there are various tools available that provide infrastructure to package, deploy and serve these models. Two of the popular ones are BentoML and MLFlow.

BentoML

BentoML is a flexible framework for building and deploying production-ready machine learning services. It allows data scientists to packaging their trained models, their dependencies, and the infrastructure code required to serve the model into a reusable package called a “Bento”.

BentoML supports various machine learning frameworks and deployment platforms and provides a unified API for managing the lifecycle of the model. Once a model is packaged as a Bento, it can be deployed to various serving platforms like AWS Lambda, Kubernetes, or Docker. BentoML also offers an API server that can be used to serve the model via a REST API. You can know more about it here. 

MLFlow

MLFlow is an open-source platform for managing the end-to-end machine learning lifecycle. It provides a comprehensive set of tools for tracking experiments, packaging code, and dependencies, and deploying models. 

MLFlow allows data scientists to easily package their models in a standard format that can be deployed to various platforms like AWS SageMaker, Azure ML, and Google Cloud AI Platform. The platform also provides a model registry to manage model versions and track their performance over time. Additionally, MLFlow offers a REST API for serving models, which can be easily integrated into web applications or other services.

How to store ML models?

Now that we know about saving models let’s see how we can store them to facilitate their quick and easy retrieval.

Storing ML models in a database

There is also scope for you to save your ML models in relational databases PostgreSQL, MySQL, Oracle SQL, etc. or NoSQL databases like MongoDB, Cassandra, etc. The choice of database totally depends on factors such as the type and volume of data being stored, the performance and scalability requirements, and the specific needs of the application. 

PostgreSQL is a popular choice when working on ML models that provide support for storing and manipulating structured data. Storing ML models in PostgreSQL provides an easy way to keep track of different versions of a model and manage them in a centralized location. 

Additionally, it allows for easy sharing of models across a team or organization. However, it’s important to note that storing large models in a database can increase database size and query times, so it’s important to consider the storage capacity and performance of your database when storing models in PostgreSQL.

To save an ML model in a database like PostgreSQL, you need to first Convert the trained model into a serialized format, such as a byte stream (pickle object) or JSON.

import pickle

# serialize the model
model_bytes = pickle.dumps(model)

Then open a connection to the database and create a table or collection to store the serialized model. For this, you need to use the psycopg2 library of Python, which lets you connect to the PostgreSQL database. You can download this library with the help of the Python package installer. 

$ pip install psycopg2-binary

Then you need to establish a connection to the database to store the ML model like this:

import psycopg2

#  establishing the connection to the Database
conn = psycopg2.connect(
  database="database-name", user=user-name, password='your-password', host='127.0.0.1', port= '5432'
)

To perform any operation on the database, you need to create a cursor object that will help you to execute queries in your Python program.

# create a cursor
cur = conn.cursor()

With the help of this cursor, you can now execute the CREATE TABLE query to create a new table.

cur.execute("CREATE TABLE models (id INT PRIMARY KEY NOT NULL, name CHAR(50), model BYTEA)")

Note: Make sure that the model object type is BYTEA. 

Finally, you can store the model and other metadata information using the INSERT INTO command. 

# Insert the serialized model into the database
cur.execute("INSERT INTO models (id, name, model) VALUES (%s, %s, %s)", (1, 'iris-classifier', model_bytes))
conn.commit()

# Close the database connection
cur.close()
conn.close()

Once all the operations are done, close the cursor and connection to the database. 

Finally, to read the model from the database, you can use the SELECT command by filtering the model either on name or id. 

import psycopg2
import pickle

# Connect to the database
conn = psycopg2.connect(
  database="database-name", user=user-name, password='your-password', host='127.0.0.1', port= '5432'
)

# Retrieve the serialized model from the database
cur = conn.cursor()
cur.execute("SELECT model FROM models WHERE name = %s", ('iris-classifier',))
model_bytes = cur.fetchone()[0]

# Deserialize the model
model = pickle.loads(model_bytes)

# Close the database connection
cur.close()
conn.close()

Once the model is loaded from the database, you can use it to make predictions as follows:

# test loaded model
y_predict = model.predict(X_test)

# check results
print(classification_report(y_test, y_predict)) 

This is it, you have the model stored and loaded from the database. 

Pros of storing ML models in a database 

Cons of storing ML models in a database

While pickle, joblib, and JSON are common ways to save machine learning models, they have limitations when it comes to versioning, sharing, and managing machine learning models. Here ML model registries come to the rescue and resolve all the issues faced by the alternatives. 

Next, you will see how saving ML models in the model registry can help you achieve reproducibility and reusability. 

Storing ML models in model registry

• A model registry is a central repository that can store, version, and manage machine learning models. 
• It typically includes features like model versioning, metadata control, comparing model runs, etc. 
• When working on any ML or DL projects, you can save and retrieve the models and their metadata from the model registry anytime you want. 
• Above all, model registries enable high collaboration among team members. 

Check out this article to learn more about model registry.  

There are various options for the model registry, for example, neptune.ai, Mlflow, Kubeflow, etc. Although all these platforms have some unique features on their own, it is rather wise to choose a registry that can provide you with a wide set of features. 

In this example, I’ll use the Neptune. It has a model registry functionality developed for organizing, storing, and managing machine learning models. It’s a great option for data scientists and ML engineers that need to manage their trained models, because it gives them collaboration features, user-friendly interface, and model versioning capabilities. 

You can set up a free account here or learn more about the tool here.

Register a model to Neptune registry

Once you have created a free account,  you can click on the New Project button to start a new project. 

Once the project creation is done, you will see a page with different configurations for saving the model. With Neptune, you can work with different frameworks like Scikit-Learn, Keras, Tensorflow, Pytorch, and more.

To store models in the Neptune model registry, you need to install the library:

$ pip install neptune

Note: Make sure that you have dumped your trained model in a file using the pickle or joblib module to store it in the model registry. 

Once the dependency is installed you need to import it into your program and initialize the Neptune model by providing it a name, a unique key (in capital letters), and your Neptune credentials. You can find all of this information in the model metadata tab of a Neptune project.

import neptune

# initialize model
model = neptune.init_model(
name="Prediction model",
key="IRMOD", # should be different for your project
project="username/project-name",
api_token="your-api-key", # your credentials
)

In the above code, Neptune dependency is imported and a model (which you want to store and track with Neptune’s model registry) is initialized with the Neptune credentials. Then you need to assign the classification model metadata to the Neptune model object. 

# assign the classification model metadata to model object
model_info = {"size_limit": 7.09, "size_units": "KB"}
model["model"] = model_info

Finally, you can upload the model to the Neptune model registry using the upload() method like this:

# upload the model to registry
model["model/signature"].upload("iris_classifier_model.pkl")

Additionally, you can track the dataset version using the track_files() method provided by neptune. 

# track dataset version
model["data/train_and_test"].track_files("iris.data")

# top the session
model.stop()

This is it, your model and the dataset are now saved to the registry. Also, don’t forget to close the session with the stop() method. 

Version a model

When you work on a real-world ML project, you happen to try a lot of models and a combination of parameters and hyperparameters. If you don’t keep track of this data, you might not know all things you have tried out, and there could be possible rework. 

This is where the Neptune model registry helps you, as you can register different versions of a model with only a few lines of code. To begin, you need to initialize a ModelVersion object as follows:

# initialize the ModelVersion
import neptune
model_version = neptune.init_model_version(
model="IR-IRMOD", # should be different for your project
project="username/projectname",
api_token="your-api-key", # your credentials
)

Then you can optionally save the model and other metadata details in each model version that you are going to register in the Neptune registry. 

# model parameters
parameters = {
"algorithm": clf_model.get_params()['algorithm'],
"max_iter": clf_model.get_params()['leaf_size'],
"solver": clf_model.get_params()['metric'],
"metric_params": clf_model.get_params()['metric_params'],
"n_jobs": clf_model.get_params()['n_jobs'],
"n_neighbors": clf_model.get_params()['n_neighbors'],
"p": clf_model.get_params()['p'],
"weights": clf_model.get_params()['weights'],
}

# log model parameters and other metadata
model_version["model/binary"].upload("iris_classifier_model.pkl")
model_version["model/parameters"] = parameters
model_version["data/dataset"].track_files("iris.data")
model_version["validation/acc"] = 0.93

# stop session
model_version.stop()

Once done, you can stop the session with the stop() method. 

Query model and metadata from registry 

Finally, you need to access this saved model and metadata when needed. You can load any specific model version that you have saved to the registry. For this, you need to initialize a ModelVersion object by providing it with the model version id. 

import neptune
import pickle

# specify model version id
version_id = 'IR-IRMOD-1' # should be different for your project

# initialize run
model_version = neptune.init_model_version(
with_id=version_id,
project="username/project-name",
api_token="your-api-key", # your credentials
)

Once done, you can access different model objects like the model, metadata, dataset, etc. details that you had registered. To begin with, let’s download the model from the registry and save it locally to test its performance on test data.

# save model from registry to local
if model_version.exists("model/binary"):
model_version["model/binary"].download(f"model/{version_id}_model.pkl")

# load model from saved pickle file
with open(f"model/{version_id}_model.pkl", 'rb') as file: 
clf_model_2 = pickle.load(file)

# evaluate model
y_predict = clf_model_2.predict(X_test)

You can also check the model metadata information that you have saved in Neptune. 

To download this metadata locally, you can use the following code:

# check metadata
model_version["model/parameters"].fetch()

This is it, you now know about storing and loading specific models from the Neptune model registry.

You can read more about the Neptune model registry here.

Pros of storing models with model registry 

Cons of storing models with model registry 

You have now seen different ways of saving an ML model (model registry being the most optimal one), this is time to check some ways to save the Deep Learning (DL) based models. 

Best practices

In this section, you will see some of the best practices for saving the ML and DL models. 

• Ensure Library Versions: Using different library versions for saving and loading the models may create compatibility issues as there could be some structural changes with the library update. You must ensure that library versions while loading the machine learning models should be the same as the library versions used to save the model. 
• Ensure Python Versions: It is a good practice to use the same Python version across all stages of your ML pipeline development. Sometimes changes in the Python version can create execution issues, for example, TensorflowV1 is supported up till Python 3.7, and if you try to use it with later versions, you will face the errors. 
• Save Both Model Architecture and Weights: In the case of DL-based models, if you save only model weight but not architecture, then you can not reconstruct the model. Saving the model architecture along with the trained weights ensures that the model can be fully reconstructed and used later on.
• Document the Model: The goal, inputs, outputs, and anticipated performance of the model should be documented. This can aid others in understanding the capabilities and constraints of the model.
• Use Model Registry: Use a model registry like neptune.ai to keep track of models, their versions, and metadata and to collaborate with team members. 
• Keep the Saved Model Secure: Keep the saved model secure by encrypting it or storing it in a secure location, especially if it contains sensitive data.

Conclusions

In conclusion, saving machine learning models is an important step in the development process, as it allows you to reuse and share your models with others. There are several ways to save machine learning models, each with its own advantages and disadvantages. Some popular methods include using pickle, Joblib, JSON, TensorFlow save, and PyTorch save.

It is important to choose the appropriate file format for your specific use case and to follow best practices for saving and documenting models, such as version control, ensuring language and library versions, and testing the saved model. By following the practices discussed in this article, you can ensure that your machine-learning models are saved correctly, are easy to reuse and deploy, and can be effectively shared with others.

References

1. https://machinelearningmastery.com/save-load-machine-learning-models-python-scikit-learn/  
2. https://www.tensorflow.org/tutorials/keras/save_and_load 
3. https://pytorch.org/tutorials/beginner/saving_loading_models.html 
4. https://www.kaggle.com/code/prmohanty/python-how-to-save-and-load-ml-models

Was the article useful?

Yes No

Thank you for your feedback!

Thanks for your vote! It's been noted. | What topics you would like to see for your next read?

Thanks for your vote! It's been noted. | Let us know what should be improved.

Your message (optional)

Submit

Δ

Thanks! Your suggestions have been forwarded to our editors

More about How to Save Trained Model in Python

Check out our product resources and related articles below:

Related article

Building MLOps Capabilities at GitLab As a One-Person ML Platform Team

Read more Related article

How to Optimize Hyperparameter Search Using Bayesian Optimization and Optuna

Read more Related article

Customizing LLM Output: Post-Processing Techniques

Read more Related article

Deep Learning Optimization Algorithms

Read more

Explore more content topics:

Computer Vision General LLMOps ML Model Development ML Tools MLOps Natural Language Processing Product Updates Reinforcement Learning Tabular Data Time Series