{ "cells": [ { "cell_type": "markdown", "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": { "byteLimit": 2048000, "rowLimit": 10000 }, "inputWidgets": {}, "nuid": "7066087c-6947-4a50-a7fb-6ab71c004ccf", "showTitle": false, "title": "" } }, "source": [ "# Log Your PyTorch Model to mlflow\n", "\n", "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.\n", "\n", "We will cover how to:\n", "- Define a simple pytorch model\n", "- Set a model signature for our logged model to define inputs and outputs to the mlflow model\n", "- Log our model to MLflow server\n", "- Load the model back from storage to use in other notebooks" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Start mlflow Server\n", "You can either:\n", "- Start a local tracking server by running `mlflow server` within the same directory that your notebook is in\n", " - Please follow [this section of the contributing guide](https://github.com/mlflow/mlflow/blob/master/CONTRIBUTING.md#javascript-and-ui) to get the UI set up.\n", "- Use a tracking server, as described in [this overview](https://mlflow.org/docs/latest/getting-started/tracking-server-overview/index.html)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Install dependencies" ] }, { "cell_type": "code", "execution_count": 1, "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": { "byteLimit": 2048000, "rowLimit": 10000 }, "inputWidgets": {}, "nuid": "885e8968-cc18-48db-af92-1b034c6ea3af", "showTitle": false, "title": "" }, "jupyter": { "outputs_hidden": true } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "\n", "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m23.2.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m23.3.1\u001b[0m\n", "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n", "Note: you may need to restart the kernel to use updated packages.\n" ] } ], "source": [ "%pip install -q mlflow torch torchmetrics" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Import packages" ] }, { "cell_type": "code", "execution_count": 2, "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": { "byteLimit": 2048000, "rowLimit": 10000 }, "inputWidgets": {}, "nuid": "092ae65d-cebd-4778-8db5-400097966a3e", "showTitle": false, "title": "" } }, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "/Users/bryan.qiu/.pyenv/versions/3.8.13/envs/mlflow/lib/python3.8/site-packages/pydantic/_internal/_fields.py:149: UserWarning: Field \"model_server_url\" has conflict with protected namespace \"model_\".\n", "\n", "You may be able to resolve this warning by setting `model_config['protected_namespaces'] = ()`.\n", " warnings.warn(\n", "/Users/bryan.qiu/.pyenv/versions/3.8.13/envs/mlflow/lib/python3.8/site-packages/pydantic/_internal/_config.py:318: UserWarning: Valid config keys have changed in V2:\n", "* 'schema_extra' has been renamed to 'json_schema_extra'\n", " warnings.warn(message, UserWarning)\n" ] } ], "source": [ "import torch\n", "import torchmetrics\n", "from sklearn.datasets import load_iris\n", "from sklearn.model_selection import train_test_split\n", "from sklearn.preprocessing import StandardScaler\n", "from torch import nn\n", "from torch.utils.data import DataLoader, TensorDataset\n", "\n", "import mlflow\n", "import mlflow.pytorch" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Prepare the data\n", "\n", "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](https://archive.ics.uci.edu/dataset/53/iris).\n", "\n", " 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." ] }, { "cell_type": "code", "execution_count": 3, "metadata": {}, "outputs": [], "source": [ "# Load and preprocess the Iris dataset\n", "iris = load_iris()\n", "X = iris.data\n", "y = iris.target\n", "\n", "# Standardize features\n", "scaler = StandardScaler()\n", "X_scaled = scaler.fit_transform(X)\n", "\n", "# Split the data into training and testing sets\n", "X_train, X_test, y_train, y_test = train_test_split(X_scaled, y, test_size=0.2, random_state=42)\n", "\n", "# Convert arrays to PyTorch tensors\n", "X_train_tensor = torch.tensor(X_train, dtype=torch.float32)\n", "y_train_tensor = torch.tensor(y_train, dtype=torch.long)\n", "X_test_tensor = torch.tensor(X_test, dtype=torch.float32)\n", "y_test_tensor = torch.tensor(y_test, dtype=torch.long)\n", "\n", "# Create datasets and dataloaders\n", "train_dataset = TensorDataset(X_train_tensor, y_train_tensor)\n", "train_loader = DataLoader(dataset=train_dataset, batch_size=16)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Define your pytorch model" ] }, { "cell_type": "code", "execution_count": 4, "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": { "byteLimit": 2048000, "rowLimit": 10000 }, "inputWidgets": {}, "nuid": "51760571-74ad-4053-9191-b82fb061c3f5", "showTitle": false, "title": "" } }, "outputs": [], "source": [ "# Define a simple neural network model\n", "class SimpleNN(nn.Module):\n", " def __init__(self):\n", " super().__init__()\n", " self.fc1 = nn.Linear(4, 10)\n", " self.fc2 = nn.Linear(10, 3)\n", "\n", " def forward(self, x):\n", " x = torch.relu(self.fc1(x))\n", " x = self.fc2(x)\n", " return x\n", "\n", "\n", "model = SimpleNN()\n", "loss = nn.CrossEntropyLoss()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Define the model signature\n", "\n", "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`.\n", "\n", "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).\n", "\n", "For more information about infer_signature, please read [the `mlflow.models.infer_signature` docs](https://mlflow.org/docs/latest/python_api/mlflow.models.html#mlflow.models.infer_signature).\n", "\n", "If you've already logged a model, you can add a signature to the logged model with [this API](https://www.mlflow.org/docs/2.8.0/models.html#set-signature-on-logged-model) as well." ] }, { "cell_type": "code", "execution_count": 5, "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": { "byteLimit": 2048000, "rowLimit": 10000 }, "inputWidgets": {}, "nuid": "2b17c594-aa9a-4e0d-86fe-28d687dfa658", "showTitle": false, "title": "" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Model signature: inputs: \n", " [Tensor('float32', (-1, 4))]\n", "outputs: \n", " [Tensor('float32', (-1, 3))]\n", "params: \n", " None\n", "\n" ] } ], "source": [ "from mlflow.models.signature import infer_signature\n", "\n", "# Infer the signature of the model\n", "sample_input = X_train_tensor[:1]\n", "model.eval()\n", "with torch.no_grad():\n", " sample_output = model(sample_input)\n", "signature = infer_signature(sample_input.numpy(), sample_output.numpy())\n", "\n", "print(\"Model signature:\", signature)" ] }, { "cell_type": "markdown", "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": { "byteLimit": 2048000, "rowLimit": 10000 }, "inputWidgets": {}, "nuid": "4525800f-012c-4203-84d2-cb4017dc3d93", "showTitle": false, "title": "" } }, "source": [ "Start an mlflow run and see how our model performs!" ] }, { "cell_type": "code", "execution_count": 6, "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": { "byteLimit": 2048000, "rowLimit": 10000 }, "inputWidgets": {}, "nuid": "92645e70-05dd-47ea-a459-d4ca3726b77c", "showTitle": false, "title": "" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Epoch 1, Loss: 0.9870926588773727, Accuracy: 0.3359375\n", "Epoch 2, Loss: 0.9870926588773727, Accuracy: 0.3359375\n", "Epoch 3, Loss: 0.9870926588773727, Accuracy: 0.3359375\n", "Epoch 4, Loss: 0.9870926588773727, Accuracy: 0.3359375\n", "Epoch 5, Loss: 0.9870926588773727, Accuracy: 0.3359375\n", "Model training and logging complete.\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "/Users/bryan.qiu/.pyenv/versions/3.8.13/envs/mlflow/lib/python3.8/site-packages/_distutils_hack/__init__.py:33: UserWarning: Setuptools is replacing distutils.\n", " warnings.warn(\"Setuptools is replacing distutils.\")\n" ] } ], "source": [ "mlflow.set_experiment(\"iris_classification_pytorch\")\n", "\n", "# Start an MLflow run\n", "with mlflow.start_run() as run:\n", " accuracy_metric = torchmetrics.Accuracy(\n", " task=\"multiclass\", num_classes=3\n", " ) # Instantiate the Accuracy metric\n", "\n", " for epoch in range(5): # number of epochs\n", " total_loss = 0\n", " total_accuracy = 0\n", "\n", " for inputs, labels in train_loader:\n", " outputs = model(inputs)\n", " curr_loss = loss(outputs, labels)\n", " curr_loss.backward()\n", "\n", " total_loss += curr_loss.item()\n", "\n", " # Calculate accuracy using torchmetrics\n", " _, preds = torch.max(outputs, 1)\n", " total_accuracy += accuracy_metric(preds, labels).item()\n", "\n", " avg_loss = total_loss / len(train_loader)\n", " avg_accuracy = total_accuracy / len(train_loader)\n", "\n", " print(f\"Epoch {epoch + 1}, Loss: {avg_loss}, Accuracy: {avg_accuracy}\")\n", " mlflow.log_metric(\"loss\", avg_loss, step=epoch)\n", " mlflow.log_metric(\"accuracy\", avg_accuracy, step=epoch)\n", "\n", " # Log the PyTorch model with the signature\n", " mlflow.pytorch.log_model(model, name=\"model\", signature=signature)\n", "\n", " # Log parameters\n", " mlflow.log_param(\"epochs\", 10)\n", " mlflow.log_param(\"batch_size\", 16)\n", " mlflow.log_param(\"learning_rate\", 0.001)\n", "\n", "print(\"Model training and logging complete.\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Loading the logged model back into memory with `mlflow.pytorch.load_model`" ] }, { "cell_type": "code", "execution_count": 7, "metadata": { "application/vnd.databricks.v1+cell": { "cellMetadata": { "byteLimit": 2048000, "rowLimit": 10000 }, "inputWidgets": {}, "nuid": "caac8221-e7bb-4772-9946-acd4c4f331e6", "showTitle": false, "title": "" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Run id from the run above: 24cb360323474df7b9090db92237a1e0\n", "Original model output: tensor([[-0.2564, 0.4631, 0.2051]])\n", "Loaded model output: tensor([[-0.2564, 0.4631, 0.2051]])\n" ] } ], "source": [ "print(\"Run id from the run above:\", run.info.run_id)\n", "\n", "# Later, or in a different script, you can load the model using the run ID\n", "loaded_model = mlflow.pytorch.load_model(f\"runs:/{run.info.run_id}/model\")\n", "\n", "# you can now use the loaded model as you would've used the original pytorch model!\n", "loaded_model.eval()\n", "with torch.no_grad():\n", " sample_input = X_test_tensor[:1]\n", " loaded_output = loaded_model(sample_input)\n", " og_output = model(sample_input)\n", " print(\"Original model output:\", og_output)\n", " print(\"Loaded model output:\", loaded_output)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "\n", "# What you see in the mlflow UI\n", "This is what you would see on the tracking server (either local or hosted, depending on your choice at the beginning)\n", "\n", "### Experiment page\n", "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.\n", "\n", "\n", "\n", "### Runs detail page\n", "Here, you can see the run ID of this run (used to retrieve the logged model) and the model signature that we set above.\n", "\n", "\n", "\n" ] } ], "metadata": { "application/vnd.databricks.v1+notebook": { "dashboards": [], "language": "python", "notebookMetadata": { "pythonIndentUnit": 2 }, "notebookName": "pytorch log model", "widgets": {} }, "kernelspec": { "display_name": "mlflow", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.8.13" } }, "nbformat": 4, "nbformat_minor": 0 }