model-registry-201/Logging_Models_PyTorch_Lightning.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/wandb/edu/blob/main/model-registry-201/Logging_Models_PyTorch_Lightning.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>\n",
    "<!--- @wandbcode{modelreg201-pytorch-lightning-colab} -->"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"https://wandb.me/logo-im-png\" width=\"400\" alt=\"Weights & Biases\" />\n",
    "\n",
    "<!--- @wandbcode{wandb201-pytorch-lightning-colab} -->\n",
    "\n",
    "## Model Registry Tutorial\n",
    "The model registry is a central place to house and organize all the model tasks and their associated artifacts being worked on across an org:\n",
    "- Model checkpoint management\n",
    "- Document your models with rich model cards\n",
    "- Maintain a history of all the models being used/deployed\n",
    "- Facilitate clean hand-offs and stage management of models\n",
    "- Tag and organize various model tasks\n",
    "- Set up automatic notifications when models progress\n",
    "\n",
    "This tutorial will walkthrough how to track the model development lifecycle for a simple image classification task.\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 🛠️ Install `wandb`\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -q wandb onnx pytorch-lightning"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Login to W&B\n",
    "- You can explicitly login using `wandb login` or `wandb.login()` (See below)\n",
    "- Alternatively you can set environment variables. There are several env variables which you can set to change the behavior of W&B logging. The most important are:\n",
    "    - `WANDB_API_KEY` - find this in your \"Settings\" section under your profile\n",
    "    - `WANDB_BASE_URL` - this is the url of the W&B server\n",
    "- Find your API Token in \"Profile\" -> \"Setttings\" in the W&B App\n",
    "\n",
    "![api_token](https://drive.google.com/uc?export=view&id=1Xn7hnn0rfPu_EW0A_-32oCXqDmpA0-kx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import wandb\n",
    "wandb.login()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Log Data and Model Checkpoints as Artifacts  \n",
    "W&B Artifacts allows you to track and version arbitrary serialized data (e.g. datasets, model checkpoints, evaluation results). When you create an artifact, you give it a name and a type, and that artifact is forever linked to the experimental system of record. If the underlying data changes, and you log that data asset again, W&B will automatically create new versions through checksummming its contents. W&B Artifacts can be thought of as a lightweight abstraction layer on top of shared unstructured file systems.\n",
    "\n",
    "### Anatomy of an artifact\n",
    "\n",
    "The `Artifact` class will correspond to an entry in the W&B Artifact registry.  The artifact has\n",
    "* a name\n",
    "* a type\n",
    "* metadata\n",
    "* description\n",
    "* files, directory of files, or references\n",
    "\n",
    "Example usage:\n",
    "```\n",
    "run = wandb.init(project = \"my-project\")\n",
    "artifact = wandb.Artifact(name = \"my_artifact\", type = \"data\")\n",
    "artifact.add_file(\"/path/to/my/file.txt\")\n",
    "run.log_artifact(artifact)\n",
    "run.finish()\n",
    "```\n",
    "\n",
    "In this tutorial, the first thing we will do is download a training dataset and log it as an artifact to be used downstream in the training job."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#@title Enter your W&B project and entity\n",
    "\n",
    "# FORM VARIABLES\n",
    "PROJECT_NAME = \"model-registry-201\" #@param {type:\"string\"}\n",
    "ENTITY = \"wandb\"#@param {type:\"string\"}\n",
    "\n",
    "src_url = \"https://storage.googleapis.com/wandb_datasets/nature_100.zip\"\n",
    "src_zip = \"nature_100.zip\"\n",
    "DATA_SRC = \"nature_100\"\n",
    "IMAGES_PER_LABEL = 10\n",
    "BALANCED_SPLITS = {\"train\" : 8, \"val\" : 1, \"test\": 1}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "!curl -SL $src_url > $src_zip\n",
    "!unzip $src_zip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import wandb\n",
    "import pandas as pd\n",
    "import os\n",
    "\n",
    "with wandb.init(project=PROJECT_NAME, entity=ENTITY, job_type='log_datasets') as run:\n",
    "  img_paths = []\n",
    "  for root, dirs, files in os.walk('nature_100', topdown=False):\n",
    "    for name in files:\n",
    "        img_path = os.path.join(root, name)\n",
    "        label = img_path.split('/')[1]\n",
    "        img_paths.append([img_path, label])\n",
    "\n",
    "  index_df = pd.DataFrame(columns=['image_path', 'label'], data=img_paths)\n",
    "  index_df.to_csv('index.csv', index=False)\n",
    "\n",
    "  train_art = wandb.Artifact(name='Nature_100', type='raw_images', description='nature image dataset with 10 classes, 10 images per class')\n",
    "  train_art.add_dir('nature_100')\n",
    "\n",
    "  # Also adding a csv indicating the labels of each image\n",
    "  train_art.add_file('index.csv')\n",
    "  wandb.log_artifact(train_art)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Using Artifact names and aliases to easily hand-off and abstract data assets\n",
    "- By simply referring to the `name:alias` combination of a dataset or model, we can better standardize components of a workflow\n",
    "- For instance, you can build PyTorch `Dataset`'s or `DataModule`'s which take as arguments W&B Artifact names and aliases to load appropriately\n",
    "\n",
    "You can now see all the metadata associated with this dataset, the W&B runs consuming it, and the whole lineage of upstream and downstream artifacts!\n",
    "\n",
    "![api_token](https://drive.google.com/uc?export=view&id=1fEEddXMkabgcgusja0g8zMz8whlP2Y5P)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchvision import transforms\n",
    "import pytorch_lightning as pl\n",
    "import torch\n",
    "from torch.utils.data import Dataset, DataLoader, random_split\n",
    "from skimage import io, transform\n",
    "from torchvision import transforms, utils, models\n",
    "import math\n",
    "\n",
    "class NatureDataset(Dataset):\n",
    "    def __init__(self,\n",
    "                 wandb_run,\n",
    "                 artifact_name_alias=\"Nature_100:latest\",\n",
    "                 local_target_dir=\"Nature_100:latest\",\n",
    "                 transform=None):\n",
    "        self.local_target_dir = local_target_dir\n",
    "        self.transform = transform\n",
    "\n",
    "        # Pull down the artifact locally to load it into memory\n",
    "        art = wandb_run.use_artifact(artifact_name_alias)\n",
    "        path_at = art.download(root=self.local_target_dir)\n",
    "\n",
    "        self.ref_df = pd.read_csv(os.path.join(self.local_target_dir, 'index.csv'))\n",
    "        self.class_names = self.ref_df.iloc[:, 1].unique().tolist()\n",
    "        self.idx_to_class = {k: v for k, v in enumerate(self.class_names)}\n",
    "        self.class_to_idx = {v: k for k, v in enumerate(self.class_names)}\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.ref_df)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        if torch.is_tensor(idx):\n",
    "            idx = idx.tolist()\n",
    "\n",
    "        img_path = self.ref_df.iloc[idx, 0]\n",
    "\n",
    "        image = io.imread(img_path)\n",
    "        label = self.ref_df.iloc[idx, 1]\n",
    "        label = torch.tensor(self.class_to_idx[label], dtype=torch.long)\n",
    "\n",
    "        if self.transform:\n",
    "            image = self.transform(image)\n",
    "\n",
    "        return image, label\n",
    "\n",
    "\n",
    "class NatureDatasetModule(pl.LightningDataModule):\n",
    "    def __init__(self,\n",
    "                 wandb_run,\n",
    "                 artifact_name_alias: str = \"Nature_100:latest\",\n",
    "                 local_target_dir: str = \"Nature_100:latest\",\n",
    "                 batch_size: int = 16,\n",
    "                 input_size: int = 224,\n",
    "                 seed: int = 42):\n",
    "        super().__init__()\n",
    "        self.wandb_run = wandb_run\n",
    "        self.artifact_name_alias = artifact_name_alias\n",
    "        self.local_target_dir = local_target_dir\n",
    "        self.batch_size = batch_size\n",
    "        self.input_size = input_size\n",
    "        self.seed = seed\n",
    "\n",
    "    def setup(self, stage=None):\n",
    "        self.nature_dataset = NatureDataset(wandb_run=self.wandb_run,\n",
    "                                            artifact_name_alias=self.artifact_name_alias,\n",
    "                                            local_target_dir=self.local_target_dir,\n",
    "                                            transform=transforms.Compose([transforms.ToTensor(),\n",
    "                                                                          transforms.CenterCrop(self.input_size),\n",
    "                                                                          transforms.Normalize((0.485, 0.456, 0.406),\n",
    "                                                                                               (0.229, 0.224, 0.225))]))\n",
    "\n",
    "        nature_length = len(self.nature_dataset)\n",
    "        train_size = math.floor(0.8 * nature_length)\n",
    "        val_size = math.floor(0.2 * nature_length)\n",
    "        self.nature_train, self.nature_val = random_split(self.nature_dataset,\n",
    "                                                          [train_size, val_size],\n",
    "                                                          generator=torch.Generator().manual_seed(self.seed))\n",
    "        return self\n",
    "\n",
    "    def train_dataloader(self):\n",
    "        return DataLoader(self.nature_train, batch_size=self.batch_size)\n",
    "\n",
    "    def val_dataloader(self):\n",
    "        return DataLoader(self.nature_val, batch_size=self.batch_size)\n",
    "\n",
    "    def predict_dataloader(self):\n",
    "        pass\n",
    "\n",
    "    def teardown(self, stage: str):\n",
    "        pass"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model Training"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Writing the Model Class and Validation Function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import wandb\n",
    "from torch.nn import Linear, CrossEntropyLoss, functional as F\n",
    "from torch.optim import Adam\n",
    "from torchmetrics.functional import accuracy\n",
    "from pytorch_lightning import LightningModule\n",
    "from torchvision import models\n",
    "import pandas as pd\n",
    "import os\n",
    "from pytorch_lightning.callbacks import ModelCheckpoint\n",
    "from pytorch_lightning.loggers import WandbLogger\n",
    "from pytorch_lightning import Trainer\n",
    "from pytorch_lightning.callbacks import Callback\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from torch.optim.lr_scheduler import StepLR\n",
    "\n",
    "class LogPredictionsCallback(Callback):\n",
    "    def __init__(self):\n",
    "      super().__init__()\n",
    "\n",
    "    def on_validation_epoch_start(self, trainer, pl_module):\n",
    "      self.batch_dfs = []\n",
    "      self.image_list = []\n",
    "      self.val_table = wandb.Table(columns=['image', 'ground_truth', 'prediction'])\n",
    "\n",
    "\n",
    "    def on_validation_batch_end(\n",
    "      self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx):\n",
    "      \"\"\"Called when the validation batch ends.\"\"\"\n",
    "\n",
    "      # Append validation predictions and ground truth to log in confusion matrix\n",
    "      x, y = batch\n",
    "      preds, y = outputs\n",
    "      self.batch_dfs.append(pd.DataFrame({\"Ground Truth\": y.numpy(), \"Predictions\": preds.numpy()}))\n",
    "\n",
    "      # Add wandb.Image to a table to log at the end of validation\n",
    "      x = x.numpy().transpose(0, 2, 3, 1)\n",
    "      for x_i, y_i, y_pred in list(zip(x, y, preds)):\n",
    "        self.image_list.append(wandb.Image(x_i, caption=f'Ground Truth: {y_i} - Prediction: {y_pred}'))\n",
    "        self.val_table.add_data(wandb.Image(x_i), y_i, y_pred)\n",
    "\n",
    "\n",
    "    def on_validation_epoch_end(self, trainer, pl_module):\n",
    "      # Collect statistics for whole validation set and log\n",
    "      class_names = trainer.datamodule.nature_dataset.class_names\n",
    "      val_df = pd.concat(self.batch_dfs)\n",
    "      wandb.log({\"eval/val_table\": self.val_table,\n",
    "                 \"eval/images_over_time\": self.image_list,\n",
    "                 \"eval/conf_matrix\": wandb.plot.confusion_matrix(y_true = val_df[\"Ground Truth\"].tolist(),\n",
    "                                                                       preds=val_df[\"Predictions\"].tolist(),\n",
    "                                                                       class_names=class_names)}, step=trainer.global_step)\n",
    "\n",
    "      del self.batch_dfs\n",
    "      del self.val_table\n",
    "\n",
    "def set_parameter_requires_grad(model, feature_extracting):\n",
    "    if feature_extracting:\n",
    "        for param in model.parameters():\n",
    "            param.requires_grad = False\n",
    "\n",
    "\n",
    "def initialize_model(model_name, num_classes, feature_extract, use_pretrained=True):\n",
    "    # Initialize these variables which will be set in this if statement. Each of these\n",
    "    #   variables is model specific.\n",
    "    model_ft = None\n",
    "    input_size = 0\n",
    "\n",
    "    if model_name == \"squeezenet\":\n",
    "        \"\"\" Squeezenet\n",
    "        \"\"\"\n",
    "        model_ft = models.squeezenet1_0(pretrained=use_pretrained)\n",
    "        set_parameter_requires_grad(model_ft, feature_extract)\n",
    "        model_ft.classifier[1] = torch.nn.Conv2d(512, num_classes, kernel_size=(1, 1), stride=(1, 1))\n",
    "        model_ft.num_classes = num_classes\n",
    "        input_size = 224\n",
    "\n",
    "    else:\n",
    "        print(\"Invalid model name, exiting...\")\n",
    "        exit()\n",
    "\n",
    "    return model_ft, input_size\n",
    "\n",
    "class NatureLitModule(LightningModule):\n",
    "    def __init__(self,\n",
    "                 model_name,\n",
    "                 num_classes=10,\n",
    "                 feature_extract=True,\n",
    "                 lr=0.01):\n",
    "        '''method used to define our model parameters'''\n",
    "        super().__init__()\n",
    "\n",
    "        self.model_name = model_name\n",
    "        self.num_classes = num_classes\n",
    "        self.feature_extract = feature_extract\n",
    "        self.model, self.input_size = initialize_model(model_name=self.model_name,\n",
    "                                                       num_classes=self.num_classes,\n",
    "                                                       feature_extract=True)\n",
    "\n",
    "        # loss\n",
    "        self.loss = CrossEntropyLoss()\n",
    "\n",
    "        # optimizer parameters\n",
    "        self.lr = lr\n",
    "\n",
    "        # save hyper-parameters to self.hparams (auto-logged by W&B)\n",
    "        self.save_hyperparameters()\n",
    "\n",
    "        # Record the gradients of all the layers\n",
    "        wandb.watch(self.model)\n",
    "\n",
    "    def forward(self, x):\n",
    "        '''method used for inference input -> output'''\n",
    "        x = self.model(x)\n",
    "\n",
    "        return x\n",
    "\n",
    "    def training_step(self, batch, batch_idx):\n",
    "        '''needs to return a loss from a single batch'''\n",
    "        preds, y, loss, acc = self._get_preds_loss_accuracy(batch)\n",
    "\n",
    "        # Log loss and metric\n",
    "        self.log('train/loss', loss)\n",
    "        self.log('train/accuracy', acc)\n",
    "\n",
    "        return loss\n",
    "\n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        '''used for logging metrics'''\n",
    "        preds, y, loss, acc = self._get_preds_loss_accuracy(batch)\n",
    "\n",
    "        # Log loss and metric\n",
    "        self.log('eval/loss', loss)\n",
    "        self.log('eval/accuracy', acc)\n",
    "\n",
    "        # Let's return preds to use it in a custom callback\n",
    "        return preds, y\n",
    "\n",
    "    def test_step(self, batch, batch_idx):\n",
    "        '''used for logging metrics'''\n",
    "        preds, y, loss, acc = self._get_preds_loss_accuracy(batch)\n",
    "\n",
    "        # Log loss and metric\n",
    "        self.log('test/loss', loss)\n",
    "        self.log('test/accuracy', acc)\n",
    "\n",
    "    def configure_optimizers(self):\n",
    "        '''defines model optimizer'''\n",
    "        return Adam(self.parameters(), lr=self.lr)\n",
    "\n",
    "\n",
    "    def _get_preds_loss_accuracy(self, batch):\n",
    "        '''convenience function since train/valid/test steps are similar'''\n",
    "        x, y = batch\n",
    "        logits = self(x)\n",
    "        preds = torch.argmax(logits, dim=1)\n",
    "        loss = self.loss(logits, y)\n",
    "\n",
    "        acc = accuracy(preds, y, task=\"multiclass\", num_classes=self.num_classes)\n",
    "        return preds, y, loss, acc"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Log Checkpoints Automatically with PyTorch Lightning\n",
    "\n",
    "Use PyTorch Lightning's `ModelCheckpoint` callback to automatically log checkpoints to W&B\n",
    "\n",
    "```\n",
    "wandb_logger = WandbLogger(log_model='all', checkpoint_name=f'image-segmentation-{wandb.run.id}')\n",
    "```\n",
    "\n",
    "- Log all models if `log_model=\"all\"` or at end of training if `log_model=True`.\n",
    "- Optionally use the wandb artifacts api to implement your own checkpointing logic using PyTorch Lightning Callbacks\n",
    "\n",
    "See more details on our PyTorch Lightning integration [here](https://docs.wandb.ai/guides/integrations/lightning)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wandb.init(project=PROJECT_NAME,\n",
    "            entity=ENTITY,\n",
    "            job_type='training',\n",
    "            config={'train_dataset': \"Nature_100:latest\",\n",
    "                    \"local_dataset_dir\": \"nature_100\",\n",
    "                    'model_name': 'squeezenet',\n",
    "                    'lr': 1.0,\n",
    "                    'gamma': 0.75,\n",
    "                    'batch_size': 16,\n",
    "                    'epochs': 5})\n",
    "\n",
    "wandb.config['input_size'] = 224\n",
    "\n",
    "wandb_logger = WandbLogger(log_model='all', checkpoint_name=f'image-segmentation-{wandb.run.id}')\n",
    "\n",
    "checkpoint_callback = ModelCheckpoint(every_n_epochs=1, save_top_k=5, monitor=\"eval/accuracy\")\n",
    "\n",
    "model = NatureLitModule(model_name=wandb.config['model_name']) # Access hyperparameters downstream to instantiate models/datasets\n",
    "\n",
    "nature_module = NatureDatasetModule(wandb_run = wandb_logger.experiment,\n",
    "                                    artifact_name_alias = wandb.config[\"train_dataset\"],\n",
    "                                    local_target_dir = wandb.config[\"local_dataset_dir\"],\n",
    "                                    batch_size=wandb.config['batch_size'],\n",
    "                                    input_size=model.input_size)\n",
    "nature_module.setup()\n",
    "\n",
    "trainer = Trainer(logger=wandb_logger,  # W&B integration\n",
    "                  accelerator=\"auto\",\n",
    "                  callbacks=[checkpoint_callback],\n",
    "                  max_epochs=5,\n",
    "                  log_every_n_steps=5)\n",
    "trainer.fit(model, datamodule=nature_module)\n",
    "\n",
    "wandb.finish()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Manage all your model checkpoints for a project under one roof.\n",
    "\n",
    "![api_token](https://drive.google.com/uc?export=view&id=1z7nXRgqHTPYjfR1SoP-CkezyxklbAZlM)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Note: Syncing with W&B Offline\n",
    "If for some reason, network communication is lost during the course of training, you can always sync progress with `wandb sync`\n",
    "\n",
    "The W&B sdk caches all logged data in a local directory `wandb` and when you call `wandb sync`, this syncs the your local state with the web app."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model Registry\n",
    "After logging a bunch of checkpoints across multiple runs during experimentation, now comes time to hand-off the best checkpoint to the next stage of the workflow (e.g. testing, deployment).\n",
    "\n",
    "The Model Registry is a central page that lives above individual W&B projects. It houses **Registered Models**, portfolios that store \"links\" to the valuable checkpoints living in individual W&B Projects.\n",
    "\n",
    "The model registry offers a centralized place to house the best checkpoints for all your model tasks. Any `model` artifact you log can be \"linked\" to a Registered Model."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Creating **Registered Models** and Linking through the UI\n",
    "#### 1. Access your team's model registry by going the team page and selecting `Model Registry`\n",
    "\n",
    "![model registry](https://drive.google.com/uc?export=view&id=1ZtJwBsFWPTm4Sg5w8vHhRpvDSeQPwsKw)\n",
    "\n",
    "#### 2. Create a new Registered Model.\n",
    "\n",
    "![model registry](https://drive.google.com/uc?export=view&id=1RuayTZHNE0LJCxt1t0l6-2zjwiV4aDXe)\n",
    "\n",
    "#### 3. Go to the artifacts tab of the project that holds all your model checkpoints\n",
    "\n",
    "![model registry](https://drive.google.com/uc?export=view&id=1LfTLrRNpBBPaUb_RmBIE7fWFMG0h3e0E)\n",
    "\n",
    "#### 4. Click \"Link to Registry\" for the model artifact version you want."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Creating Registered Models and Linking through the **API**\n",
    "You can [link a model via api](https://docs.wandb.ai/guides/models) with `wandb.run.link_artifact` passing in the artifact object, and the name of the **Registered Model**, along with aliases you want to append to it. **Registered Models** are entity (team) scoped in W&B so only members of a team can see and access the **Registered Models** there. You indicate a registered model name via api with `<entity>/model-registry/<registered-model-name>`. If a Registered Model doesn't exist, one will be created automatically."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "last_run_id = \"rmlp8vlj\" #@param\n",
    "wandb.init(project=PROJECT_NAME, entity=ENTITY, job_type=\"registering_best_model\")\n",
    "best_model = wandb.use_artifact(f'{ENTITY}/{PROJECT_NAME}/image-segmentation-{last_run_id}:latest')\n",
    "registered_model_name = \"YOLOv8 Image Segmentation\" #@param {type: \"string\"}\n",
    "wandb.run.link_artifact(best_model, f'{ENTITY}/model-registry/{registered_model_name}', aliases=['staging'])\n",
    "wandb.finish()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### What is \"Linking\"?\n",
    "When you link to the registry, this creates a new version of that Registered Model, which is just a pointer to the artifact version living in that project. There's a reason W&B segregates the versioning of artifacts in a project from the versioning of a Registered Model. The process of linking a model artifact version is equivalent to \"bookmarking\" that artifact version under a Registered Model task.\n",
    "\n",
    "Typically during R&D/experimentation, researchers generate 100s, if not 1000s of model checkpoint artifacts, but only one or two of them actually \"see the light of day.\" This process of linking those checkpoints to a separate, versioned registry helps delineate the model development side from the model deployment/consumption side of the workflow. The globally understood version/alias of a model should be unpolluted from all the experimental versions being generated in R&D and thus the versioning of a Registered Model increments according to new \"bookmarked\" models as opposed to model checkpoint logging."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create a Centralized Hub for all your models\n",
    "- Add a model card, tags, slack notifactions to your Registered Model\n",
    "- Change aliases to reflect when models move through different phases\n",
    "- Embed the model registry in reports for model documentation and regression reports. See this report as an [example](https://api.wandb.ai/links/wandb-smle/r82bj9at)\n",
    "![model registry](https://drive.google.com/uc?export=view&id=1lKPgaw-Ak4WK_91aBMcLvUMJL6pDQpgO)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set up Slack Notifications when new models get linked to the registry\n",
    "\n",
    "![model registry](https://drive.google.com/uc?export=view&id=1RsWCa6maJYD5y34gQ0nwWiKSWUCqcjT9)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Consuming a Registered Model\n",
    "You now can consume any registered model via API by referring the corresponding `name:alias`. Model consumers, whether they are engineers, researchers, or CI/CD processes, can go to the model registry as the central hub for all models that should \"see the light of day\": those that need to go through testing or move to production."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%wandb -h 600\n",
    "\n",
    "run = wandb.init(project=PROJECT_NAME, entity=ENTITY, job_type='inference')\n",
    "artifact = run.use_artifact(f'{ENTITY}/model-registry/Model Registry Tutorial:staging', type='model')\n",
    "artifact_dir = artifact.download()\n",
    "wandb.finish()"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "include_colab_link": true,
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}