integrations/torchvision/tutorials/docs-torchvision-sparsify-from-scratch-resnet50-beans.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "00e081b7",
   "metadata": {},
   "source": [
    "# Sparsifying ResNet-50 from Scratch (Beans)\n",
    "\n",
    "In this example, we will demonstrate how to sparsify an image classification model from scratch using SparseML's PyTorch integration. We train and prune [ResNet-50](https://pytorch.org/vision/main/models/generated/torchvision.models.resnet50.html) on the downstream [Beans dataset](https://huggingface.co/datasets/beans) using the Global Magnitude Pruning algorithm. \n",
    "\n",
    "## Agenda\n",
    "\n",
    "There are a few steps:\n",
    "\n",
    " 1. Setup the dataset\n",
    " 2. Setup the PyTorch training loop\n",
    " 3. Train a dense version of ResNet-50\n",
    " 4. Run the GMP pruning algorithm on the dense model\n",
    " \n",
    "## Installation\n",
    "\n",
    "Install SparseML and `datasets` with `pip`:\n",
    "\n",
    "```\n",
    "pip install sparseml[torchvision]\n",
    "pip install datasets\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "1ad80edf",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import sparseml\n",
    "import torchvision\n",
    "from sparseml.pytorch.optim import ScheduledModifierManager\n",
    "from sparseml.pytorch.utils import TensorBoardLogger, ModuleExporter, get_prunable_layers, tensor_sparsity\n",
    "from torch.utils.data import DataLoader\n",
    "from torch.nn import CrossEntropyLoss\n",
    "from torch.optim import Adam\n",
    "from torchvision import transforms\n",
    "from tqdm.auto import tqdm\n",
    "import math\n",
    "import datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c421a838",
   "metadata": {},
   "source": [
    "## Step 1: Setup Dataset\n",
    "\n",
    "Beans leaf dataset is a set of images of diseased and healthy leaves. Based on a leaf image, the goal of this task is to predict the disease type (Angular Leaf Spot and Bean Rust), if any.\n",
    "\n",
    "We will use the Hugging Face `datasets` library to download the data and the torchvision `ImageFolder` in the training loop.\n",
    "\n",
    "[Checkout the dataset card](https://huggingface.co/datasets/beans)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f49303e",
   "metadata": {},
   "outputs": [],
   "source": [
    "beans_dataset = datasets.load_dataset(\"beans\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "4115c9d8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/ubuntu/.cache/huggingface/datasets/downloads/extracted/7ad2d437b751e134577576a32849c44a9ade89297680ad5f6a64051e2108810b/train/angular_leaf_spot/angular_leaf_spot_train.0.jpg\n",
      "/home/ubuntu/.cache/huggingface/datasets/downloads/extracted/1141f86479bc0bb56c75616d153591cc8299d1ea4edc53bb1ab65edd2c65b240/validation/angular_leaf_spot/angular_leaf_spot_val.0.jpg\n"
     ]
    }
   ],
   "source": [
    "print(beans_dataset[\"train\"][0][\"image_file_path\"])\n",
    "print(beans_dataset[\"validation\"][0][\"image_file_path\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "00aaf463",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_path = \"/home/ubuntu/.cache/huggingface/datasets/downloads/extracted/7ad2d437b751e134577576a32849c44a9ade89297680ad5f6a64051e2108810b/train\"\n",
    "val_path = \"/home/ubuntu/.cache/huggingface/datasets/downloads/extracted/1141f86479bc0bb56c75616d153591cc8299d1ea4edc53bb1ab65edd2c65b240/validation\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "7d180bbc",
   "metadata": {},
   "outputs": [],
   "source": [
    "NUM_LABELS = 3\n",
    "BATCH_SIZE = 64\n",
    "\n",
    "# imagenet transforms\n",
    "imagenet_transform = transforms.Compose([\n",
    "   transforms.Resize(size=256, interpolation=transforms.InterpolationMode.BILINEAR, max_size=None, antialias=None),\n",
    "   transforms.CenterCrop(size=(224, 224)),\n",
    "   transforms.ToTensor(),\n",
    "   transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
    "])\n",
    "\n",
    "# datasets\n",
    "train_dataset = torchvision.datasets.ImageFolder(\n",
    "    root=train_path,\n",
    "    transform=imagenet_transform\n",
    ")\n",
    "\n",
    "val_dataset = torchvision.datasets.ImageFolder(\n",
    "    root=val_path,\n",
    "    transform=imagenet_transform\n",
    ")\n",
    "\n",
    "\n",
    "# dataloaders\n",
    "train_loader = DataLoader(train_dataset, BATCH_SIZE, shuffle=True, pin_memory=True, num_workers=16)\n",
    "val_loader = DataLoader(val_dataset, BATCH_SIZE, shuffle=False, pin_memory=True, num_workers=16)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "132fc514",
   "metadata": {},
   "source": [
    "## Step 2: Setup PyTorch Training Loop\n",
    "\n",
    "We will use this training loop below. This is standard PyTorch functionality."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "5f1b878a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cuda\n"
     ]
    }
   ],
   "source": [
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "print(device)\n",
    "\n",
    "def run_model_one_epoch(model, data_loader, criterion, device, train=False, optimizer=None):\n",
    "    if train:\n",
    "        model.train()\n",
    "    else:\n",
    "        model.eval()\n",
    "\n",
    "    running_loss = 0.0\n",
    "    total_correct = 0\n",
    "    total_predictions = 0\n",
    "\n",
    "    # loop through batches\n",
    "    for step, (inputs, labels) in tqdm(enumerate(data_loader), total=len(data_loader)):\n",
    "        inputs = inputs.to(device)\n",
    "        labels = labels.to(device)\n",
    "\n",
    "        if train:\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "        # compute loss, run backpropogation\n",
    "        outputs = model(inputs)  # model returns logits\n",
    "        loss = criterion(outputs, labels)\n",
    "        if train:\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "        running_loss += loss.item()\n",
    "\n",
    "        # run evaluation\n",
    "        predictions = outputs.argmax(dim=1)\n",
    "        total_correct += torch.sum(predictions == labels).item()\n",
    "        total_predictions += inputs.size(0)\n",
    "\n",
    "    # return loss and evaluation metric\n",
    "    loss = running_loss / (step + 1.0)\n",
    "    accuracy = total_correct / total_predictions\n",
    "    return loss, accuracy"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21e250c6",
   "metadata": {},
   "source": [
    "## **Step 3: Train ResNet-50 on Beans**\n",
    "\n",
    "First, we will train a dense version of ResNet-50 on the Beans dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "4d554578",
   "metadata": {},
   "outputs": [],
   "source": [
    "# download pre-trained model, setup classification head\n",
    "model = torchvision.models.resnet50(weights=torchvision.models.ResNet50_Weights.DEFAULT)\n",
    "model.fc = torch.nn.Linear(model.fc.in_features, NUM_LABELS)\n",
    "model.to(device)\n",
    "\n",
    "# setup loss function and optimizer\n",
    "criterion = CrossEntropyLoss()\n",
    "optimizer = Adam(model.parameters(), lr=8e-3) # lr will be override by sparseml"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39a82a98",
   "metadata": {},
   "source": [
    "Next, we will use SparseML's recipes to set the hyperparameters of training loop. In this case, we will use the following recipe:\n",
    "\n",
    "```yaml\n",
    "# Epoch and Learning-Rate variables\n",
    "num_epochs: 10.0\n",
    "init_lr: 0.0005\n",
    "\n",
    "training_modifiers:\n",
    "  - !EpochRangeModifier\n",
    "    start_epoch: 0.0\n",
    "    end_epoch: eval(num_epochs)\n",
    "\n",
    "  - !LearningRateFunctionModifier\n",
    "    final_lr: 0.0\n",
    "    init_lr: eval(init_lr)\n",
    "    lr_func: cosine\n",
    "    start_epoch: 0.0\n",
    "    end_epoch: eval(num_epochs)\n",
    "```\n",
    "\n",
    "As you can see, the recipe includes an `!EpochRangeModifier` and a `!LearningRateFunctionModifier`. These modifiers simply set the number of epochs to train for and the learning rate schedule. As a result, the final model will be dense."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4553beb",
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat ./recipes/resnet50-beans-dense-recipe.yaml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "1b19a8d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "dense_recipe_path = \"./recipes/resnet50-beans-dense-recipe.yaml\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e00a7a9",
   "metadata": {},
   "source": [
    "Next, we use SparseML's `ScheduledModifierManager` to parse and apply the recipe. The `manager.modify` function modifies and wraps the `model` and `optimizer` with the instructions from the recipe. You can use the `model` and `optimizer` just like standard PyTorch objects."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "1749e00d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# create ScheduledModifierManager and Optimizer wrapper\n",
    "manager = ScheduledModifierManager.from_yaml(dense_recipe_path)\n",
    "optimizer = manager.modify(model, optimizer, steps_per_epoch=len(train_loader))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f3b56352",
   "metadata": {},
   "source": [
    "Kick off the transfer learning loop. Our run reached ~99% validation accuracy after 10 epochs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "d00d175b",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running Training Epoch 1/10\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d5b96520239d4d728d23a4f099eb8fd3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/17 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training Epoch: 1/10\n",
      "Training Loss: 0.45206236138063316\n",
      "Top 1 Acc: 0.8085106382978723\n",
      "\n",
      "Running Validation Epoch 1/10\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "287feebeb9514f33b58e0cb6a825f672",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation Epoch: 1/10\n",
      "Val Loss: 0.8609340712428093\n",
      "Top 1 Acc: 0.849624060150376\n",
      "\n",
      "Running Training Epoch 2/10\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "76336e81f4c24749b17ca07b2ce5bd69",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/17 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training Epoch: 2/10\n",
      "Training Loss: 0.1085617478717776\n",
      "Top 1 Acc: 0.9584139264990329\n",
      "\n",
      "Running Validation Epoch 2/10\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c34d773bae5c48ce96743bd2b807a6f7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation Epoch: 2/10\n",
      "Val Loss: 0.1460044514387846\n",
      "Top 1 Acc: 0.9699248120300752\n",
      "\n",
      "Running Training Epoch 3/10\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "aa655dafb1514e9684fd3d949effd615",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/17 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training Epoch: 3/10\n",
      "Training Loss: 0.07007143213687574\n",
      "Top 1 Acc: 0.97678916827853\n",
      "\n",
      "Running Validation Epoch 3/10\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d4c8ea7aa9d04edea382c7f14a03a065",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation Epoch: 3/10\n",
      "Val Loss: 0.027990046694564324\n",
      "Top 1 Acc: 0.9849624060150376\n",
      "\n",
      "Running Training Epoch 4/10\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "34da8fe5ac4a48d584517b9861c6965b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/17 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training Epoch: 4/10\n",
      "Training Loss: 0.017229604699155864\n",
      "Top 1 Acc: 0.9970986460348162\n",
      "\n",
      "Running Validation Epoch 4/10\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1f06331bc3a3451f8220574fa627d139",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation Epoch: 4/10\n",
      "Val Loss: 0.024875935167074203\n",
      "Top 1 Acc: 1.0\n",
      "\n",
      "Running Training Epoch 5/10\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ecbbd1253c8c474d9822ca9aa3f32358",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/17 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training Epoch: 5/10\n",
      "Training Loss: 0.0032389334005796734\n",
      "Top 1 Acc: 1.0\n",
      "\n",
      "Running Validation Epoch 5/10\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9a3cd1a358764830b5b657cc77f2c57c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation Epoch: 5/10\n",
      "Val Loss: 0.020098081634690363\n",
      "Top 1 Acc: 0.9849624060150376\n",
      "\n",
      "Running Training Epoch 6/10\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e454370d9ac948c9ac2015297a578ff6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/17 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training Epoch: 6/10\n",
      "Training Loss: 0.0009401099643368713\n",
      "Top 1 Acc: 1.0\n",
      "\n",
      "Running Validation Epoch 6/10\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e881bba812114c39a09b3341efc5a4ae",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation Epoch: 6/10\n",
      "Val Loss: 0.013677676247122387\n",
      "Top 1 Acc: 0.9924812030075187\n",
      "\n",
      "Running Training Epoch 7/10\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1bdc58cad336471580d42b9c869f8379",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/17 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training Epoch: 7/10\n",
      "Training Loss: 0.0007901448304491008\n",
      "Top 1 Acc: 1.0\n",
      "\n",
      "Running Validation Epoch 7/10\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "071da6885f4644f4abb9e3c68850ea07",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation Epoch: 7/10\n",
      "Val Loss: 0.011986067402176559\n",
      "Top 1 Acc: 1.0\n",
      "\n",
      "Running Training Epoch 8/10\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "94495ac2721a4112b6b53b42ac16787c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/17 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training Epoch: 8/10\n",
      "Training Loss: 0.001196134924505125\n",
      "Top 1 Acc: 1.0\n",
      "\n",
      "Running Validation Epoch 8/10\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c2a2315f371443c89e8492cadb9b98ae",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation Epoch: 8/10\n",
      "Val Loss: 0.008087300811894238\n",
      "Top 1 Acc: 1.0\n",
      "\n",
      "Running Training Epoch 9/10\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e4cbaef71d224dcab1918a8e68738567",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/17 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training Epoch: 9/10\n",
      "Training Loss: 0.0012429110587725196\n",
      "Top 1 Acc: 1.0\n",
      "\n",
      "Running Validation Epoch 9/10\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "458f09dd3c004d65b683786268d2bce1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation Epoch: 9/10\n",
      "Val Loss: 0.01152486567540715\n",
      "Top 1 Acc: 1.0\n",
      "\n",
      "Running Training Epoch 10/10\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bd1541379cdd496b8fd93b8ced597067",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/17 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training Epoch: 10/10\n",
      "Training Loss: 0.017825346116082715\n",
      "Top 1 Acc: 0.9990328820116054\n",
      "\n",
      "Running Validation Epoch 10/10\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c952c716ba6d46d9919d29b05fd5b2df",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation Epoch: 10/10\n",
      "Val Loss: 0.01295137139580523\n",
      "Top 1 Acc: 0.9849624060150376\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# run transfer learning\n",
    "epoch = 0\n",
    "for epoch in range(manager.max_epochs):\n",
    "    # run training loop\n",
    "    epoch_name = f\"{epoch + 1}/{manager.max_epochs}\"\n",
    "    print(f\"Running Training Epoch {epoch_name}\")\n",
    "    train_loss, train_acc = run_model_one_epoch(model, train_loader, criterion, device, train=True, optimizer=optimizer)\n",
    "    print(f\"Training Epoch: {epoch_name}\\nTraining Loss: {train_loss}\\nTop 1 Acc: {train_acc}\\n\")\n",
    "\n",
    "    # run validation loop\n",
    "    print(f\"Running Validation Epoch {epoch_name}\")\n",
    "    val_loss, val_acc = run_model_one_epoch(model, val_loader, criterion, device)\n",
    "    print(f\"Validation Epoch: {epoch_name}\\nVal Loss: {val_loss}\\nTop 1 Acc: {val_acc}\\n\")\n",
    "\n",
    "manager.finalize(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "11fa4cf2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sparseml.pytorch.utils import ModuleExporter\n",
    "\n",
    "save_dir = \"dense_model\"\n",
    "exporter = ModuleExporter(model, output_dir=save_dir)\n",
    "exporter.export_pytorch(name=\"resnet-50-dense-beans.pth\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "1617f476",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "848572dd",
   "metadata": {},
   "source": [
    "## Step 4: Prune The Model\n",
    "\n",
    "With a model trained on Beans, we are now ready to apply the GMP algorithm to prune the model. The GMP algorithm is an interative pruning algorithm. At the end of each epoch, we identify the lowest magnitude weights (those closest to 0) and remove them from the network starting from an initial level of sparsity until a final level of sparsity. The remaining nonzero weights are then fine-tuned onto training dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "8447d074",
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint = torch.load(\"./dense_model/training/resnet-50-dense-beans.pth\")\n",
    "model = torchvision.models.resnet50()\n",
    "model.fc = torch.nn.Linear(model.fc.in_features, NUM_LABELS)\n",
    "model.load_state_dict(checkpoint['state_dict'])\n",
    "model.to(device)\n",
    "\n",
    "# setup loss function and optimizer, LR will be overriden by sparseml\n",
    "criterion = CrossEntropyLoss()\n",
    "optimizer = Adam(model.parameters(), lr=8e-3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "edb62dc9",
   "metadata": {
    "scrolled": false
   },
   "source": [
    "Next, we need to create a SparseML recipe which includes the GMP algorithm. The `!GlobalMagnitudePruningModifier` modifier instructs SparseML to apply the GMP algorithm at a global level (pruning the lowest magnitude weights across all layers).\n",
    "\n",
    "Firstly, we need to decide identify which parameters of the model to apply the GMP algorithm to. We can use the `get_prunable_layers` function to inspect:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "8950185f",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "conv1\n",
      "layer1.0.conv1\n",
      "layer1.0.conv2\n",
      "layer1.0.conv3\n",
      "layer1.0.downsample.0\n",
      "layer1.1.conv1\n",
      "layer1.1.conv2\n",
      "layer1.1.conv3\n",
      "layer1.2.conv1\n",
      "layer1.2.conv2\n",
      "layer1.2.conv3\n",
      "layer2.0.conv1\n",
      "layer2.0.conv2\n",
      "layer2.0.conv3\n",
      "layer2.0.downsample.0\n",
      "layer2.1.conv1\n",
      "layer2.1.conv2\n",
      "layer2.1.conv3\n",
      "layer2.2.conv1\n",
      "layer2.2.conv2\n",
      "layer2.2.conv3\n",
      "layer2.3.conv1\n",
      "layer2.3.conv2\n",
      "layer2.3.conv3\n",
      "layer3.0.conv1\n",
      "layer3.0.conv2\n",
      "layer3.0.conv3\n",
      "layer3.0.downsample.0\n",
      "layer3.1.conv1\n",
      "layer3.1.conv2\n",
      "layer3.1.conv3\n",
      "layer3.2.conv1\n",
      "layer3.2.conv2\n",
      "layer3.2.conv3\n",
      "layer3.3.conv1\n",
      "layer3.3.conv2\n",
      "layer3.3.conv3\n",
      "layer3.4.conv1\n",
      "layer3.4.conv2\n",
      "layer3.4.conv3\n",
      "layer3.5.conv1\n",
      "layer3.5.conv2\n",
      "layer3.5.conv3\n",
      "layer4.0.conv1\n",
      "layer4.0.conv2\n",
      "layer4.0.conv3\n",
      "layer4.0.downsample.0\n",
      "layer4.1.conv1\n",
      "layer4.1.conv2\n",
      "layer4.1.conv3\n",
      "layer4.2.conv1\n",
      "layer4.2.conv2\n",
      "layer4.2.conv3\n",
      "fc\n"
     ]
    }
   ],
   "source": [
    "# print parameters\n",
    "for (name, layer) in get_prunable_layers(model):\n",
    "    print(f\"{name}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5ef3bdf",
   "metadata": {},
   "source": [
    "We will apply pruning to each of the convs and exclude the fc layer (which is the final projection head). Fortunately, SparseML allows us to pass regexes to identify layers in the network, so we can use the following list to identify the relevant layers for pruning:\n",
    "    \n",
    "    - 'conv1.weight'\n",
    "    - 're:layer1.*.conv1.weight'\n",
    "    - 're:layer1.*.conv2.weight'\n",
    "    - 're:layer1.*.conv3.weight'\n",
    "    - 're:layer1.0.downsample.0.weight'\n",
    "    - 're:layer2.*.conv1.weight'\n",
    "    - 're:layer2.*.conv2.weight'\n",
    "    - 're:layer2.*.conv3.weight'\n",
    "    - 're:layer2.0.downsample.0.weight'\n",
    "    - 're:layer3.*.conv1.weight'\n",
    "    - 're:layer3.*.conv2.weight'\n",
    "    - 're:layer3.*.conv3.weight'\n",
    "    - 're:layer3.0.downsample.0.weight'\n",
    "    - 're:layer4.*.conv1.weight'\n",
    "    - 're:layer4.*.conv2.weight'\n",
    "    - 're:layer4.*.conv3.weight'\n",
    "    - 're:layer4.0.downsample.0.weight'\n",
    "\n",
    "Here is what the recipe looks like:\n",
    "\n",
    "```yaml\n",
    "# Epoch hyperparams\n",
    "stabilization_epochs: 1.0\n",
    "pruning_epochs: 9.0\n",
    "finetuning_epochs: 5.0\n",
    "\n",
    "# Learning rate hyperparams\n",
    "init_lr: 0.0001\n",
    "final_lr: 0.00005\n",
    "\n",
    "# Pruning hyperparams\n",
    "init_sparsity: 0.05\n",
    "final_sparsity: 0.9\n",
    "\n",
    "# Stabalization Stage\n",
    "training_modifiers:\n",
    "  - !EpochRangeModifier\n",
    "    start_epoch: 0.0\n",
    "    end_epoch: eval(stabilization_epochs + pruning_epochs + finetuning_epochs)\n",
    "  \n",
    "  - !SetLearningRateModifier\n",
    "    start_epoch: 0.0\n",
    "    learning_rate: eval(init_lr)\n",
    "\n",
    "# Pruning Stage\n",
    "pruning_modifiers:\n",
    "  - !LearningRateFunctionModifier\n",
    "    init_lr: eval(init_lr)\n",
    "    final_lr: eval(final_lr)\n",
    "    lr_func: cosine\n",
    "    start_epoch: eval(stabilization_epochs)\n",
    "    end_epoch: eval(stabilization_epochs + pruning_epochs)\n",
    "    \n",
    "  - !GlobalMagnitudePruningModifier\n",
    "    init_sparsity: eval(init_sparsity)\n",
    "    final_sparsity: eval(final_sparsity)\n",
    "    start_epoch: eval(stabilization_epochs)\n",
    "    end_epoch: eval(stabilization_epochs + pruning_epochs)\n",
    "    update_frequency: 0.5\n",
    "    params:        \n",
    "        - 'conv1.weight'\n",
    "        - 're:layer1.*.conv1.weight'\n",
    "        - 're:layer1.*.conv2.weight'\n",
    "        - 're:layer1.*.conv3.weight'\n",
    "        - 're:layer1.0.downsample.0.weight'\n",
    "        - 're:layer2.*.conv1.weight'\n",
    "        - 're:layer2.*.conv2.weight'\n",
    "        - 're:layer2.*.conv3.weight'\n",
    "        - 're:layer2.0.downsample.0.weight'\n",
    "        - 're:layer3.*.conv1.weight'\n",
    "        - 're:layer3.*.conv2.weight'\n",
    "        - 're:layer3.*.conv3.weight'\n",
    "        - 're:layer3.0.downsample.0.weight'\n",
    "        - 're:layer4.*.conv1.weight'\n",
    "        - 're:layer4.*.conv2.weight'\n",
    "        - 're:layer4.*.conv3.weight'\n",
    "        - 're:layer4.0.downsample.0.weight'\n",
    "    leave_enabled: True\n",
    "\n",
    "# Finetuning Stage\n",
    "finetuning_modifiers:\n",
    "  - !LearningRateFunctionModifier\n",
    "    init_lr: eval(init_lr)\n",
    "    final_lr: eval(final_lr)\n",
    "    lr_func: cosine\n",
    "    start_epoch: eval(stabilization_epochs + pruning_epochs)\n",
    "    end_epoch: eval(stabilization_epochs + pruning_epochs + finetuning_epochs) \n",
    "```\n",
    "\n",
    "This recipe specifies that we will run the GMP algorithm for 9 epochs after running one warmup epoch. We start at an init_sparsity level of 5% and gradually increase sparsity to a final_sparsity level of 90% following a cubic curve. The pruning is applied in an unstructured manner, meaning that any weight can be pruned.\n",
    "\n",
    "Over the final 5 epochs, we will fine-tune the 90% pruned model further. Since we set leave_enabled=True the sparsity level will be maintained as the fine-tuning occurs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c492e33",
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat recipes/resnet50-beans-pruning-recipe.yaml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "9cc0e8ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "pruning_recipe_path = \"./recipes/resnet50-beans-pruning-recipe.yaml\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "68646bc7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# create ScheduledModifierManager and Optimizer wrapper\n",
    "manager = ScheduledModifierManager.from_yaml(pruning_recipe_path)\n",
    "logger = TensorBoardLogger(log_path=\"./tensorboard_outputs\")\n",
    "optimizer = manager.modify(model, optimizer, loggers=[logger], steps_per_epoch=len(train_loader))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ab5b75b",
   "metadata": {},
   "source": [
    "Next, kick off the GMP training loop. \n",
    "\n",
    "As you can see, we use the wrapped `optimizer` and `model` in the same way as above. SparseML parsed the recipe and updated the `optimizer` with the logic of GMP algorithm from the recipe. This allows you to just the `optimizer` and `model` as usual, with all of the pruning-related logic specified by the declarative recipe interface.\n",
    "\n",
    "Our 90% pruned model reaches ~99% validation accuracy (vs ~99% for the dense model)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "601c8c21",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running Training Epoch 1/15\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1c7ce2fc1cfd45979ada8fa0b6fcf3e6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/17 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training Epoch: 1/15\n",
      "Training Loss: 0.00033969324495581747\n",
      "Top 1 Acc: 1.0\n",
      "\n",
      "Running Validation Epoch 1/15\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e3aef9e2789941b0ae700fde6e266376",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation Epoch: 1/15\n",
      "Val Loss: 0.013618851000501309\n",
      "Top 1 Acc: 0.9849624060150376\n",
      "\n",
      "Running Training Epoch 2/15\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fffd99128b11403f9b57011320732c16",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/17 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training Epoch: 2/15\n",
      "Training Loss: 9.199756892913309e-05\n",
      "Top 1 Acc: 1.0\n",
      "\n",
      "Running Validation Epoch 2/15\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e328d77edeb4431e90201713fdb617dd",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation Epoch: 2/15\n",
      "Val Loss: 0.025973695794164087\n",
      "Top 1 Acc: 0.9849624060150376\n",
      "\n",
      "Running Training Epoch 3/15\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2e68e294bf494959be7e44222cfca3ae",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/17 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training Epoch: 3/15\n",
      "Training Loss: 0.0006384848421235159\n",
      "Top 1 Acc: 1.0\n",
      "\n",
      "Running Validation Epoch 3/15\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "13a6a67a5ccf49468cc98169c5b425a8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation Epoch: 3/15\n",
      "Val Loss: 0.0031044411274100034\n",
      "Top 1 Acc: 1.0\n",
      "\n",
      "Running Training Epoch 4/15\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "db7f16560b2044a5b7bf4b99684ccaba",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/17 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training Epoch: 4/15\n",
      "Training Loss: 0.0006065369069641319\n",
      "Top 1 Acc: 1.0\n",
      "\n",
      "Running Validation Epoch 4/15\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a663816b9381480eb5f924efd32c8fd3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation Epoch: 4/15\n",
      "Val Loss: 0.002252104241051711\n",
      "Top 1 Acc: 1.0\n",
      "\n",
      "Running Training Epoch 5/15\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "18df31a72edb49e4a1f1c747589bdac3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/17 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training Epoch: 5/15\n",
      "Training Loss: 0.031195201322963627\n",
      "Top 1 Acc: 0.9990328820116054\n",
      "\n",
      "Running Validation Epoch 5/15\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8bc095e5df214848a718235df4717e74",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation Epoch: 5/15\n",
      "Val Loss: 0.01123013386192421\n",
      "Top 1 Acc: 1.0\n",
      "\n",
      "Running Training Epoch 6/15\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1d08639063644773b61b000c5b1f2efe",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/17 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training Epoch: 6/15\n",
      "Training Loss: 0.010756760572001119\n",
      "Top 1 Acc: 0.9990328820116054\n",
      "\n",
      "Running Validation Epoch 6/15\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "68c1be15dc7d4af68ec6c19153f692f9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation Epoch: 6/15\n",
      "Val Loss: 0.07213001201550166\n",
      "Top 1 Acc: 0.9924812030075187\n",
      "\n",
      "Running Training Epoch 7/15\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1e97320b9902466ea5c3c4eae6bdf9a2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/17 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training Epoch: 7/15\n",
      "Training Loss: 0.026524252669118783\n",
      "Top 1 Acc: 0.9970986460348162\n",
      "\n",
      "Running Validation Epoch 7/15\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3ce2e377d4dc4cc79f0f486ba85cee39",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation Epoch: 7/15\n",
      "Val Loss: 0.177118182182312\n",
      "Top 1 Acc: 0.9774436090225563\n",
      "\n",
      "Running Training Epoch 8/15\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a49d19cebfea407eac41f8f48f1026ce",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/17 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training Epoch: 8/15\n",
      "Training Loss: 0.02232004324083819\n",
      "Top 1 Acc: 0.9990328820116054\n",
      "\n",
      "Running Validation Epoch 8/15\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7a76f59cc2944c69b531baf0f07bd609",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation Epoch: 8/15\n",
      "Val Loss: 0.10358777642250061\n",
      "Top 1 Acc: 0.9699248120300752\n",
      "\n",
      "Running Training Epoch 9/15\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f5ebb79fa89345a7ae603be67cbe179b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/17 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training Epoch: 9/15\n",
      "Training Loss: 0.029385707779403997\n",
      "Top 1 Acc: 0.9990328820116054\n",
      "\n",
      "Running Validation Epoch 9/15\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "33b7f4317a3d42d0b7baa010ebc2bdcb",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation Epoch: 9/15\n",
      "Val Loss: 0.042455013220508896\n",
      "Top 1 Acc: 1.0\n",
      "\n",
      "Running Training Epoch 10/15\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "dbdb5527e3684766bc7015ff58bdce55",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/17 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training Epoch: 10/15\n",
      "Training Loss: 0.007427054580629748\n",
      "Top 1 Acc: 1.0\n",
      "\n",
      "Running Validation Epoch 10/15\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b73a0d01704c495e90d46ce5b418ae3b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation Epoch: 10/15\n",
      "Val Loss: 0.020314963534474373\n",
      "Top 1 Acc: 1.0\n",
      "\n",
      "Running Training Epoch 11/15\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d48f532ce1fe443986390a95d10a214a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/17 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training Epoch: 11/15\n",
      "Training Loss: 0.005854974952800309\n",
      "Top 1 Acc: 1.0\n",
      "\n",
      "Running Validation Epoch 11/15\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "66decd3b06e14086a29442a5dc339bce",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation Epoch: 11/15\n",
      "Val Loss: 0.013980435362706581\n",
      "Top 1 Acc: 1.0\n",
      "\n",
      "Running Training Epoch 12/15\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "761dd586eea2462d8059642a6fa01861",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/17 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training Epoch: 12/15\n",
      "Training Loss: 0.002781623138991349\n",
      "Top 1 Acc: 1.0\n",
      "\n",
      "Running Validation Epoch 12/15\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ef312ea80a14430ebfc5367acbb70d52",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation Epoch: 12/15\n",
      "Val Loss: 0.011425975399712721\n",
      "Top 1 Acc: 1.0\n",
      "\n",
      "Running Training Epoch 13/15\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "322dec36498c44a3ae953f2ea2364734",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/17 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training Epoch: 13/15\n",
      "Training Loss: 0.004056999764596934\n",
      "Top 1 Acc: 0.9990328820116054\n",
      "\n",
      "Running Validation Epoch 13/15\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7a431075ae9345cf894e5aa41a086257",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation Epoch: 13/15\n",
      "Val Loss: 0.013013879070058465\n",
      "Top 1 Acc: 1.0\n",
      "\n",
      "Running Training Epoch 14/15\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "67fa6ab896e945b2a29e27f361e6075a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/17 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training Epoch: 14/15\n",
      "Training Loss: 0.002845817831043592\n",
      "Top 1 Acc: 1.0\n",
      "\n",
      "Running Validation Epoch 14/15\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7e16741614424bb3b4ae46bc60ca99ea",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation Epoch: 14/15\n",
      "Val Loss: 0.010720587490747372\n",
      "Top 1 Acc: 1.0\n",
      "\n",
      "Running Training Epoch 15/15\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bb350d3d25914b7cbc88a2c642abbd3d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/17 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training Epoch: 15/15\n",
      "Training Loss: 0.009813380410538657\n",
      "Top 1 Acc: 0.9980657640232108\n",
      "\n",
      "Running Validation Epoch 15/15\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bc2e3959fb7945ef821b9c58291bdff6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation Epoch: 15/15\n",
      "Val Loss: 0.011411373193065325\n",
      "Top 1 Acc: 1.0\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# run GMP algorithm\n",
    "epoch = 0\n",
    "for epoch in range(manager.max_epochs):\n",
    "    # run training loop\n",
    "    epoch_name = f\"{epoch + 1}/{manager.max_epochs}\"\n",
    "    print(f\"Running Training Epoch {epoch_name}\")\n",
    "    train_loss, train_acc = run_model_one_epoch(model, train_loader, criterion, device, train=True, optimizer=optimizer)\n",
    "    print(f\"Training Epoch: {epoch_name}\\nTraining Loss: {train_loss}\\nTop 1 Acc: {train_acc}\\n\")\n",
    "\n",
    "    # run validation loop\n",
    "    print(f\"Running Validation Epoch {epoch_name}\")\n",
    "    val_loss, val_acc = run_model_one_epoch(model, val_loader, criterion, device)\n",
    "    print(f\"Validation Epoch: {epoch_name}\\nVal Loss: {val_loss}\\nTop 1 Acc: {val_acc}\\n\")\n",
    "    \n",
    "    logger.log_scalar(\"Metrics/Loss (Train)\", train_loss, epoch)\n",
    "    logger.log_scalar(\"Metrics/Accuracy (Train)\", train_acc, epoch)\n",
    "    logger.log_scalar(\"Metrics/Loss (Validation)\", val_loss, epoch)\n",
    "    logger.log_scalar(\"Metrics/Accuracy (Validation)\", val_acc, epoch)\n",
    "    \n",
    "manager.finalize(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "65fb7b12",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "conv1.weight: 0.4039\n",
      "layer1.0.conv1.weight: 0.4758\n",
      "layer1.0.conv2.weight: 0.6846\n",
      "layer1.0.conv3.weight: 0.5506\n",
      "layer1.0.downsample.0.weight: 0.5157\n",
      "layer1.1.conv1.weight: 0.5690\n",
      "layer1.1.conv2.weight: 0.6410\n",
      "layer1.1.conv3.weight: 0.5604\n",
      "layer1.2.conv1.weight: 0.5430\n",
      "layer1.2.conv2.weight: 0.6075\n",
      "layer1.2.conv3.weight: 0.6064\n",
      "layer2.0.conv1.weight: 0.5498\n",
      "layer2.0.conv2.weight: 0.7556\n",
      "layer2.0.conv3.weight: 0.6881\n",
      "layer2.0.downsample.0.weight: 0.8303\n",
      "layer2.1.conv1.weight: 0.7918\n",
      "layer2.1.conv2.weight: 0.8089\n",
      "layer2.1.conv3.weight: 0.7450\n",
      "layer2.2.conv1.weight: 0.7564\n",
      "layer2.2.conv2.weight: 0.8100\n",
      "layer2.2.conv3.weight: 0.7706\n",
      "layer2.3.conv1.weight: 0.7239\n",
      "layer2.3.conv2.weight: 0.7862\n",
      "layer2.3.conv3.weight: 0.7415\n",
      "layer3.0.conv1.weight: 0.6814\n",
      "layer3.0.conv2.weight: 0.8697\n",
      "layer3.0.conv3.weight: 0.7558\n",
      "layer3.0.downsample.0.weight: 0.9209\n",
      "layer3.1.conv1.weight: 0.8875\n",
      "layer3.1.conv2.weight: 0.9177\n",
      "layer3.1.conv3.weight: 0.8482\n",
      "layer3.2.conv1.weight: 0.9054\n",
      "layer3.2.conv2.weight: 0.9143\n",
      "layer3.2.conv3.weight: 0.8672\n",
      "layer3.3.conv1.weight: 0.8686\n",
      "layer3.3.conv2.weight: 0.9209\n",
      "layer3.3.conv3.weight: 0.8573\n",
      "layer3.4.conv1.weight: 0.8710\n",
      "layer3.4.conv2.weight: 0.9158\n",
      "layer3.4.conv3.weight: 0.8599\n",
      "layer3.5.conv1.weight: 0.8293\n",
      "layer3.5.conv2.weight: 0.9191\n",
      "layer3.5.conv3.weight: 0.8144\n",
      "layer4.0.conv1.weight: 0.7612\n",
      "layer4.0.conv2.weight: 0.9458\n",
      "layer4.0.conv3.weight: 0.8048\n",
      "layer4.0.downsample.0.weight: 0.9347\n",
      "layer4.1.conv1.weight: 0.9039\n",
      "layer4.1.conv2.weight: 0.9626\n",
      "layer4.1.conv3.weight: 0.8643\n",
      "layer4.2.conv1.weight: 0.9441\n",
      "layer4.2.conv2.weight: 0.9775\n",
      "layer4.2.conv3.weight: 0.8976\n",
      "fc.weight: 0.0000\n"
     ]
    }
   ],
   "source": [
    "for (name, layer) in get_prunable_layers(model):\n",
    "    print(f\"{name}.weight: {tensor_sparsity(layer.weight).item():.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00805b02",
   "metadata": {},
   "source": [
    "Finally, export your model to ONNX."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "417ac105",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/ubuntu/rshaw/sparseml/env/lib/python3.10/site-packages/torch/onnx/utils.py:439: UserWarning: It is recommended that constant folding be turned off ('do_constant_folding=False') when exporting the model in training-amenable mode, i.e. with 'training=TrainingMode.TRAIN' or 'training=TrainingMode.PRESERVE' (when model is in training mode). Otherwise, some learnable model parameters may not translate correctly in the exported ONNX model because constant folding mutates model parameters. Please consider turning off constant folding or setting the training=TrainingMode.EVAL.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "save_dir = \"experiment-0\"\n",
    "exporter = ModuleExporter(model, output_dir=save_dir)\n",
    "exporter.export_pytorch(name=\"resnet-50-sparse-beans.pth\")\n",
    "exporter.export_onnx(torch.randn(1, 3, 224, 224), name=\"sparse-model.onnx\", convert_qat=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7960612c",
   "metadata": {},
   "source": [
    "## Wrap Up\n",
    "\n",
    "The resulting model is is 90% sparse and achieves validation accuracy of ~100% (vs the unoptimized dense model at ~99%) without much hyperparameter search.\n",
    "\n",
    "Key hyperparameter experiments you may want to run include:\n",
    "- Learning rate\n",
    "- Learning rate schedule\n",
    "- Sparsity level\n",
    "- Number of pruning epochs\n",
    "\n",
    "DeepSparse supports speedup from pruning and quantization. To reach maximum performance, check out our examples of quantizing a model!"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}