2. Tutorial - pruning.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Automatic pruning\n",
    "\n",
    "This notebook demonstrates end2end pipeline for MobileNetV2 pruning.\n",
    "\n",
    "Our pruning process consists of calibration for pruning, least important channel selection, channel pruning and model fine-tuning.\n",
    "\n",
    "### Main chapters of this notebook:\n",
    "1. Setup the environment\n",
    "1. Prepare dataset and create dataloaders\n",
    "1. Evaluate pretrained MobileNetV2\n",
    "1. Calibrate, prune and evaluate pruned model\n",
    "1. Finetune and evaluate pruned model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup the environment\n",
    "\n",
    "First, let's set up the environment and make some common imports."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -r requirements.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'\n",
    "# You may need to uncomment and change this variable to match free GPU index\n",
    "# os.environ['CUDA_VISIBLE_DEVICES'] = '0'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Common:\n",
    "import logging\n",
    "import numpy as np\n",
    "import torch\n",
    "from fvcore.nn import FlopCountAnalysis\n",
    "from pathlib import Path\n",
    "from torch import nn\n",
    "from tqdm.auto import tqdm\n",
    "from typing import List\n",
    "from tutorial_utils.checkpoints import download_imagenette_mobilenet\n",
    "from tutorial_utils.dataset import create_imagenette_dataloaders_for_pruning\n",
    "from tutorial_utils.train import accuracy\n",
    "\n",
    "# Training:\n",
    "from torch.optim.lr_scheduler import CosineAnnealingLR\n",
    "from torch.optim.lr_scheduler import LinearLR\n",
    "from torch.optim.lr_scheduler import ChainedScheduler\n",
    "from torch.optim import RAdam\n",
    "from tutorial_utils.train import tutorial_train_loop\n",
    "\n",
    "# Pruning:\n",
    "from enot.pruning import KnapsackPruningLabelSelector\n",
    "from enot.pruning import PruningCalibrator\n",
    "from enot.pruning import prune_model\n",
    "from enot.utils.batch_norm import tune_bn_stats"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define model evaluation function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This function can evaluate both nn.Modules and executable functions.\n",
    "def eval_model(model, dataloader):\n",
    "    model.eval()\n",
    "\n",
    "    total = 0\n",
    "    total_loss = 0.0\n",
    "    total_correct = 0.0\n",
    "\n",
    "    criterion = nn.CrossEntropyLoss()\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for inputs, labels in tqdm(dataloader):\n",
    "            n = inputs.shape[0]\n",
    "\n",
    "            pred_labels = model(inputs)\n",
    "            batch_loss = criterion(pred_labels, labels)\n",
    "            batch_accuracy = accuracy(pred_labels, labels)\n",
    "\n",
    "            total += n\n",
    "            total_loss += batch_loss.item() * n\n",
    "            total_correct += batch_accuracy.item() * n\n",
    "\n",
    "    return total_loss / total, total_correct / total"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### In the following cell we setup all necessary dirs\n",
    "\n",
    "* `HOME_DIR` - experiments home directory\n",
    "* `DATASETS_DIR` - root directory for datasets (imagenette2, ...)\n",
    "* `PROJECT_DIR` - project directory to save training logs, checkpoints, ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "HOME_DIR = Path.home() / '.optimization_experiments'\n",
    "DATASETS_DIR = HOME_DIR / 'datasets'\n",
    "PROJECT_DIR = HOME_DIR / 'e2e_pruning'\n",
    "\n",
    "HOME_DIR.mkdir(exist_ok=True)\n",
    "DATASETS_DIR.mkdir(exist_ok=True)\n",
    "PROJECT_DIR.mkdir(exist_ok=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare dataset and create dataloaders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataloader, validation_dataloader = create_imagenette_dataloaders_for_pruning(\n",
    "    dataset_root_dir=DATASETS_DIR,\n",
    "    project_dir=PROJECT_DIR,\n",
    "    input_size=224,\n",
    "    batch_size=32,\n",
    "    num_workers=4,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluate pretrained MobileNetV2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchvision.models.mobilenetv2 import mobilenet_v2\n",
    "\n",
    "regular_model = mobilenet_v2(pretrained=False, num_classes=10).cuda()\n",
    "\n",
    "# Turning off FullyConnected layer dropout.\n",
    "# This is required to stabilize fine-tuning procedure.\n",
    "regular_model.classifier[0].p = 0.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint_path = PROJECT_DIR / 'e2e_imagenette_pruning.pth'\n",
    "download_imagenette_mobilenet(checkpoint_path)\n",
    "\n",
    "regular_model.load_state_dict(\n",
    "    torch.load(checkpoint_path)['model'],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "val_loss, val_accuracy = eval_model(regular_model, validation_dataloader)\n",
    "print(f'Regular (non-pruned) model: accuracy={val_accuracy:.3f}, loss={val_loss:.3f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Calibrate, prune and evaluate pruned model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, let's define ``mmac_acceleration_factor`` (desired acceliration in terms of macs) and loss function ``loss_function`` (calculates total loss for single batch of data loader)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mmac_acceleration_factor = 3  # This gives x3 macs reduction.\n",
    "loss_function = torch.nn.CrossEntropyLoss(reduction='sum')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now a function needs to be defined to measure latency in terms of macs for an arbitrary model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mmac_calculation_function(model):\n",
    "    inputs, _ = next(iter(train_dataloader))\n",
    "\n",
    "    fca = FlopCountAnalysis(\n",
    "        model=model.eval(),\n",
    "        inputs=inputs,\n",
    "    )\n",
    "    fca.unsupported_ops_warnings(False)\n",
    "    fca.uncalled_modules_warnings(False)\n",
    "\n",
    "    return fca.total() / 1e6"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Now, let's perform model calibration for pruning. Calibration finds all prunable channels in the network and estimates their importances. Accumulated pruning-related information will be stored in ``pruning_info`` variable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# This code is implemented in enot.pruning.calibrate_model_for_pruning.\n",
    "\n",
    "regular_model.eval()\n",
    "pruning_calibrator = PruningCalibrator(model=regular_model)\n",
    "with pruning_calibrator:\n",
    "    for images, labels in train_dataloader:\n",
    "        predictions = regular_model(images)\n",
    "        loss = loss_function(predictions, labels)\n",
    "        loss.backward()\n",
    "\n",
    "pruning_info = pruning_calibrator.pruning_info"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Selecting the least important channels to achieve the desired acceleration in terms of macs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "baseline_latency = mmac_calculation_function(model=regular_model)\n",
    "target_latency = baseline_latency / mmac_acceleration_factor\n",
    "\n",
    "print('baseline mmacs =', baseline_latency)\n",
    "print('target mmacs =', target_latency)\n",
    "\n",
    "label_selector = KnapsackPruningLabelSelector(\n",
    "    latency_calculation_function=mmac_calculation_function,\n",
    "    target_latency=target_latency,\n",
    "    verbose=True,\n",
    ")\n",
    "prune_labels = label_selector.select(model=regular_model, pruning_info=pruning_info)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a copy of original model and remove selected channels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "pruned_model = prune_model(\n",
    "    model=regular_model,\n",
    "    pruning_info=pruning_info,\n",
    "    prune_labels=prune_labels,\n",
    "    inplace=False,\n",
    ")\n",
    "pruned_model_latency = mmac_calculation_function(model=pruned_model)\n",
    "print('pruned model mmacs =', pruned_model_latency)\n",
    "pruned_model.cuda()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Tune batch normalization layers on train data to stabilize their running variables after pruning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "tune_bn_stats(\n",
    "    model=pruned_model,\n",
    "    dataloader=train_dataloader,\n",
    "    reset_bns=True,  # Drop old batch norm running statistics.\n",
    "    set_momentums_none=True,  # Accumulate average statistics.\n",
    "    n_steps=None,\n",
    "    epochs=1,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After pruning, pruned user model has the same structure as the original model, except that some convolutions, fully-connected layers and batch norm layers now have smaller number of channels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "val_loss, val_accuracy = eval_model(pruned_model, validation_dataloader)\n",
    "print(f'Pruned model: accuracy={val_accuracy:.3f}, loss={val_loss:.3f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Finetune and evaluate pruned model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N_EPOCHS = 5\n",
    "N_WARMUP_EPOCHS = 1\n",
    "learning_rate = 0.005\n",
    "\n",
    "# Uncomment lines below if you want to reach the best pruned model\n",
    "# performance (~93% accuracy for pruned model).\n",
    "\n",
    "# N_EPOCHS = 50  # Increase the number of model fine-tuning epochs.\n",
    "# N_WARMUP_EPOCHS = 10  # Increase the number of warmup epochs.\n",
    "# learning_rate = 0.01  # Increase learning rate\n",
    "\n",
    "len_train = len(train_dataloader)\n",
    "\n",
    "optimizer = RAdam(pruned_model.parameters(), lr=learning_rate, weight_decay=0.00004)\n",
    "main_scheduler = CosineAnnealingLR(optimizer, T_max=len_train * N_EPOCHS)\n",
    "warmup_scheduler = LinearLR(optimizer, start_factor=1e-16, total_iters=len_train * N_WARMUP_EPOCHS)\n",
    "scheduler = ChainedScheduler([main_scheduler, warmup_scheduler])\n",
    "loss_function = nn.CrossEntropyLoss().cuda()\n",
    "\n",
    "tutorial_train_loop(\n",
    "    epochs=N_EPOCHS,\n",
    "    model=pruned_model,\n",
    "    optimizer=optimizer,\n",
    "    metric_function=accuracy,\n",
    "    loss_function=loss_function,\n",
    "    train_loader=train_dataloader,\n",
    "    validation_loader=validation_dataloader,\n",
    "    scheduler=scheduler,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "val_loss, val_accuracy = eval_model(pruned_model, validation_dataloader)\n",
    "print(f'Fine-tuned pruned model: accuracy={val_accuracy:.3f}, loss={val_loss:.3f}')"
   ]
  }
 ],
 "metadata": {
  "execution": {
   "timeout": 3000
  },
  "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}