Log Your PyTorch Model to mlflow

This guide will walk you through how to save your PyTorch model to mlflow and load the saved model for inference. Saving a pretrained/finetuned model in MLflow allows you to easily share the model or deploy it to production.

We will cover how to:

Define a simple pytorch model
Set a model signature for our logged model to define inputs and outputs to the mlflow model
Log our model to MLflow server
Load the model back from storage to use in other notebooks

Start mlflow Server

You can either:

Start a local tracking server by running mlflow server within the same directory that your notebook is in
- Please follow this section of the contributing guide to get the UI set up.
Use a tracking server, as described in this overview

Install dependencies

python

%pip install -q mlflow torch torchmetrics

Import packages

python

import torch
import torchmetrics
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

import mlflow
import mlflow.pytorch

Prepare the data

The Iris dataset is a popular beginner's dataset for classification models that contains measurements of 3 species of Iris flowers. If you want, more information can be found at this link.

We are loading the data, standardizing it, splitting it into training and testing sets, converting it into the format required by PyTorch, and preparing it for efficient training in mini-batches.

python

# Load and preprocess the Iris dataset
iris = load_iris()
X = iris.data
y = iris.target

# Standardize features
scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)

# Split the data into training and testing sets
X_train, X_test, y_train, y_test = train_test_split(X_scaled, y, test_size=0.2, random_state=42)

# Convert arrays to PyTorch tensors
X_train_tensor = torch.tensor(X_train, dtype=torch.float32)
y_train_tensor = torch.tensor(y_train, dtype=torch.long)
X_test_tensor = torch.tensor(X_test, dtype=torch.float32)
y_test_tensor = torch.tensor(y_test, dtype=torch.long)

# Create datasets and dataloaders
train_dataset = TensorDataset(X_train_tensor, y_train_tensor)
train_loader = DataLoader(dataset=train_dataset, batch_size=16)

Define your pytorch model

python

# Define a simple neural network model
class SimpleNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(4, 10)
        self.fc2 = nn.Linear(10, 3)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x


model = SimpleNN()
loss = nn.CrossEntropyLoss()

Define the model signature

A model signature defines valid input, output and params schema, and is used to validate them at inference time. mlflow.models.infer_signature infers the model signature that can be passed into mlflow.pytorch.log_model.

Since pytorch model's usually operate on tensors, we need to convert both the input and output into a type compatible with mlflow.models.infer_signature. Commonly, this means converting them into a numpy.ndarray or a dictionary of numpy.ndarray (if the output is multiple tensors).

For more information about infer_signature, please read the mlflow.models.infer_signature docs.

If you've already logged a model, you can add a signature to the logged model with this API as well.

python

from mlflow.models.signature import infer_signature

# Infer the signature of the model
sample_input = X_train_tensor[:1]
model.eval()
with torch.no_grad():
    sample_output = model(sample_input)
signature = infer_signature(sample_input.numpy(), sample_output.numpy())

print("Model signature:", signature)

Start an mlflow run and see how our model performs!

python

mlflow.set_experiment("iris_classification_pytorch")

# Start an MLflow run
with mlflow.start_run() as run:
    accuracy_metric = torchmetrics.Accuracy(
        task="multiclass", num_classes=3
    )  # Instantiate the Accuracy metric

    for epoch in range(5):  # number of epochs
        total_loss = 0
        total_accuracy = 0

        for inputs, labels in train_loader:
            outputs = model(inputs)
            curr_loss = loss(outputs, labels)
            curr_loss.backward()

            total_loss += curr_loss.item()

            # Calculate accuracy using torchmetrics
            _, preds = torch.max(outputs, 1)
            total_accuracy += accuracy_metric(preds, labels).item()

        avg_loss = total_loss / len(train_loader)
        avg_accuracy = total_accuracy / len(train_loader)

        print(f"Epoch {epoch + 1}, Loss: {avg_loss}, Accuracy: {avg_accuracy}")
        mlflow.log_metric("loss", avg_loss, step=epoch)
        mlflow.log_metric("accuracy", avg_accuracy, step=epoch)

    # Log the PyTorch model with the signature
    mlflow.pytorch.log_model(model, name="model", signature=signature)

    # Log parameters
    mlflow.log_param("epochs", 10)
    mlflow.log_param("batch_size", 16)
    mlflow.log_param("learning_rate", 0.001)

print("Model training and logging complete.")

Loading the logged model back into memory with mlflow.pytorch.load_model

python

print("Run id from the run above:", run.info.run_id)

# Later, or in a different script, you can load the model using the run ID
loaded_model = mlflow.pytorch.load_model(f"runs:/{run.info.run_id}/model")

# you can now use the loaded model as you would've used the original pytorch model!
loaded_model.eval()
with torch.no_grad():
    sample_input = X_test_tensor[:1]
    loaded_output = loaded_model(sample_input)
    og_output = model(sample_input)
    print("Original model output:", og_output)
    print("Loaded model output:", loaded_output)

What you see in the mlflow UI

This is what you would see on the tracking server (either local or hosted, depending on your choice at the beginning)

Experiment page

Here, you can select the experiment you set in the code above and choose a run to view the model logged during that run. You can also see how your pytorch model has changed in accuracy / loss over different runs in the Chart tab.

Runs detail page

Here, you can see the run ID of this run (used to retrieve the logged model) and the model signature that we set above.