wehub-resource-sync
468 lines
15 KiB
Plaintext
468 lines
15 KiB
Plaintext
{
|
|
"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",
|
|
"<img src=\"https://i.imgur.com/hiolwEe.png\" style=\"width: 60%\">\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",
|
|
"<img src=\"https://i.imgur.com/gJN4f2v.png'\" style=\"width: 60%\">\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
|
|
}
|