tutorials/01.resnet18_cifar10.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f4441533-4d5a-4c9e-8906-1df24bda96ec",
   "metadata": {},
   "source": [
    "## Tutorial 1. ResNet18 on CIFAR10. \n",
    "\n",
    "\n",
    "In this tutorial, we will show \n",
    "\n",
    "- How to end-to-end train and compress a ResNet18 from scratch on CIFAR10 to get a compressed ResNet18.\n",
    "- The compressed ResNet18 achives both **high performance** and **significant FLOPs and parameters reductions** than the full model. \n",
    "- The compressed ResNet18 **reduces about 92% parameters** to achieve **92.91% accuracy** only lower than the baseline by **0.11%**.\n",
    "- More detailed new HESSO optimizer setup. (Technical report regarding HESSO will be released on the early of 2024)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8067c3a6-bf56-4c85-96fb-2d4da488680b",
   "metadata": {},
   "source": [
    "### Step 1. Create OTO instance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f2ce0399-51c7-4afc-b0f0-35bda3698825",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "OTO graph constructor\n",
      "graph build\n"
     ]
    }
   ],
   "source": [
    "import sys\n",
    "sys.path.append('..')\n",
    "from sanity_check.backends.resnet_cifar10 import resnet18_cifar10\n",
    "from only_train_once import OTO\n",
    "import torch\n",
    "\n",
    "model = resnet18_cifar10()\n",
    "dummy_input = torch.rand(1, 3, 32, 32)\n",
    "oto = OTO(model=model.cuda(), dummy_input=dummy_input.cuda())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "01639056-422d-4bfb-89db-ff9f0725e807",
   "metadata": {},
   "source": [
    "#### (Optional) Visualize the pruning dependancy graph of DNN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3e8cefb1-85cf-4833-a6ff-ee2f9270f02d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# A ResNet_zig.gv.pdf will be generated to display the depandancy graph.\n",
    "oto.visualize(view=False, out_dir='../cache')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc99fc48-fed7-41a5-ab6d-f2e5686f2003",
   "metadata": {},
   "source": [
    "### Step 2. Dataset Preparation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "57ec145d-4847-4f4b-8c12-782beb61cc61",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz to cifar10/cifar-10-python.tar.gz\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Extracting cifar10/cifar-10-python.tar.gz to cifar10\n",
      "Files already downloaded and verified\n"
     ]
    }
   ],
   "source": [
    "from torchvision.datasets import CIFAR10\n",
    "import torchvision.transforms as transforms\n",
    "\n",
    "trainset = CIFAR10(root='cifar10', train=True, download=True, transform=transforms.Compose([\n",
    "            transforms.RandomHorizontalFlip(),\n",
    "            transforms.RandomCrop(32, 4),\n",
    "            transforms.ToTensor(),\n",
    "            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])]))\n",
    "testset = CIFAR10(root='cifar10', train=False, download=True, transform=transforms.Compose([\n",
    "            transforms.ToTensor(),\n",
    "            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])]))\n",
    "\n",
    "trainloader =  torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True, num_workers=4)\n",
    "testloader = torch.utils.data.DataLoader(testset, batch_size=64, shuffle=False, num_workers=4)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65233349-da36-47bc-9af8-41a3f2d4429a",
   "metadata": {},
   "source": [
    "### Step 3. Setup HESSO optimizer\n",
    "\n",
    "The following main hyperparameters need to be taken care.\n",
    "\n",
    "- `variant`: The optimizer that is used for training the baseline full model. Currently support `sgd`, `adam` and `adamw`.\n",
    "- `lr`: The initial learning rate.\n",
    "- `weight_decay`: Weight decay as standard DNN optimization.\n",
    "- `target_group_sparsity`: The target group sparsity, typically higher group sparsity refers to more FLOPs and model size reduction, meanwhile may regress model performance more.\n",
    "- `start_pruning_steps`: The number of steps that **starts** to prune.\n",
    "- `pruning_steps`: The number of steps that **finishes** pruning (reach `target_group_sparsity`) after `start_pruning_steps`.\n",
    "- `pruning_periods`:  Incrementally produce the group sparsity equally among pruning periods.\n",
    "\n",
    "We empirically suggest `start_pruning_steps` as 1/10 of total number of training steps. `pruning_steps` until 1/4 or 1/5 of total number of training steps.",
    "The advatnages of HESSO compared to DHSPG is its explicit control over group sparsity exploration, which is typically more convenient."
    ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "48af4116-bb7e-4156-bc3d-c32c8fc6f2c0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Setup HESSO\n"
     ]
    }
   ],
   "source": [
    "optimizer = oto.hesso(\n",
    "    variant='sgd', \n",
    "    lr=0.1, \n",
    "    weight_decay=1e-4,\n",
    "    target_group_sparsity=0.7,\n",
    "    start_pruning_step=10 * len(trainloader), \n",
    "    pruning_periods=10,\n",
    "    pruning_steps=10 * len(trainloader)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3632cb7e-e1ae-460c-b415-81c158c80a20",
   "metadata": {},
   "source": [
    "### Step 4. Train ResNet18 as normal."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "ed9d98a1-a77b-4cc6-9acd-a1483c17e370",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ep: 0, loss: 1.56, norm_all:4124.34, grp_sparsity: 0.00, acc1: 0.3208, norm_import: 4124.34, norm_redund: 0.00, num_grp_import: 2880, num_grp_redund: 0\n",
      "Ep: 1, loss: 1.00, norm_all:4116.36, grp_sparsity: 0.00, acc1: 0.5533, norm_import: 4116.36, norm_redund: 0.00, num_grp_import: 2880, num_grp_redund: 0\n",
      "Ep: 2, loss: 0.77, norm_all:4107.12, grp_sparsity: 0.00, acc1: 0.5351, norm_import: 4107.12, norm_redund: 0.00, num_grp_import: 2880, num_grp_redund: 0\n",
      "Ep: 3, loss: 0.64, norm_all:4096.86, grp_sparsity: 0.00, acc1: 0.7380, norm_import: 4096.86, norm_redund: 0.00, num_grp_import: 2880, num_grp_redund: 0\n",
      "Ep: 4, loss: 0.56, norm_all:4086.17, grp_sparsity: 0.00, acc1: 0.7718, norm_import: 4086.17, norm_redund: 0.00, num_grp_import: 2880, num_grp_redund: 0\n",
      "Ep: 5, loss: 0.50, norm_all:4074.66, grp_sparsity: 0.00, acc1: 0.7511, norm_import: 4074.66, norm_redund: 0.00, num_grp_import: 2880, num_grp_redund: 0\n",
      "Ep: 6, loss: 0.45, norm_all:4062.94, grp_sparsity: 0.00, acc1: 0.6216, norm_import: 4062.94, norm_redund: 0.00, num_grp_import: 2880, num_grp_redund: 0\n",
      "Ep: 7, loss: 0.41, norm_all:4050.86, grp_sparsity: 0.00, acc1: 0.7732, norm_import: 4050.86, norm_redund: 0.00, num_grp_import: 2880, num_grp_redund: 0\n",
      "Ep: 8, loss: 0.38, norm_all:4038.76, grp_sparsity: 0.00, acc1: 0.8251, norm_import: 4038.76, norm_redund: 0.00, num_grp_import: 2880, num_grp_redund: 0\n",
      "Ep: 9, loss: 0.34, norm_all:4026.24, grp_sparsity: 0.00, acc1: 0.8507, norm_import: 4026.24, norm_redund: 0.00, num_grp_import: 2880, num_grp_redund: 0\n",
      "Ep: 10, loss: 0.33, norm_all:3795.98, grp_sparsity: 0.07, acc1: 0.7086, norm_import: 3795.98, norm_redund: 0.00, num_grp_import: 2679, num_grp_redund: 201\n",
      "Ep: 11, loss: 0.31, norm_all:3567.24, grp_sparsity: 0.14, acc1: 0.8409, norm_import: 3567.24, norm_redund: 0.00, num_grp_import: 2478, num_grp_redund: 402\n",
      "Ep: 12, loss: 0.30, norm_all:3342.20, grp_sparsity: 0.21, acc1: 0.7612, norm_import: 3342.20, norm_redund: 0.00, num_grp_import: 2277, num_grp_redund: 603\n",
      "Ep: 13, loss: 0.30, norm_all:3117.74, grp_sparsity: 0.28, acc1: 0.8356, norm_import: 3117.74, norm_redund: 0.00, num_grp_import: 2076, num_grp_redund: 804\n",
      "Ep: 14, loss: 0.30, norm_all:2860.77, grp_sparsity: 0.35, acc1: 0.8481, norm_import: 2860.77, norm_redund: 0.00, num_grp_import: 1875, num_grp_redund: 1005\n",
      "Ep: 15, loss: 0.29, norm_all:2563.71, grp_sparsity: 0.42, acc1: 0.8458, norm_import: 2563.71, norm_redund: 0.00, num_grp_import: 1674, num_grp_redund: 1206\n",
      "Ep: 16, loss: 0.29, norm_all:2250.46, grp_sparsity: 0.49, acc1: 0.8373, norm_import: 2250.46, norm_redund: 0.00, num_grp_import: 1473, num_grp_redund: 1407\n",
      "Ep: 17, loss: 0.30, norm_all:1957.45, grp_sparsity: 0.56, acc1: 0.8204, norm_import: 1957.45, norm_redund: 0.00, num_grp_import: 1272, num_grp_redund: 1608\n",
      "Ep: 18, loss: 0.31, norm_all:1663.43, grp_sparsity: 0.63, acc1: 0.6788, norm_import: 1663.43, norm_redund: 0.00, num_grp_import: 1071, num_grp_redund: 1809\n",
      "Ep: 19, loss: 0.35, norm_all:1362.16, grp_sparsity: 0.70, acc1: 0.7222, norm_import: 1362.16, norm_redund: 0.00, num_grp_import: 865, num_grp_redund: 2015\n",
      "Ep: 20, loss: 0.36, norm_all:1375.29, grp_sparsity: 0.70, acc1: 0.7984, norm_import: 1375.29, norm_redund: 0.00, num_grp_import: 865, num_grp_redund: 2015\n",
      "Ep: 21, loss: 0.31, norm_all:1384.57, grp_sparsity: 0.70, acc1: 0.8452, norm_import: 1384.57, norm_redund: 0.00, num_grp_import: 865, num_grp_redund: 2015\n",
      "Ep: 22, loss: 0.28, norm_all:1392.36, grp_sparsity: 0.70, acc1: 0.8700, norm_import: 1392.36, norm_redund: 0.00, num_grp_import: 865, num_grp_redund: 2015\n",
      "Ep: 23, loss: 0.27, norm_all:1399.04, grp_sparsity: 0.70, acc1: 0.8357, norm_import: 1399.04, norm_redund: 0.00, num_grp_import: 865, num_grp_redund: 2015\n",
      "Ep: 24, loss: 0.25, norm_all:1404.36, grp_sparsity: 0.70, acc1: 0.8431, norm_import: 1404.36, norm_redund: 0.00, num_grp_import: 865, num_grp_redund: 2015\n",
      "Ep: 25, loss: 0.23, norm_all:1409.66, grp_sparsity: 0.70, acc1: 0.8884, norm_import: 1409.66, norm_redund: 0.00, num_grp_import: 865, num_grp_redund: 2015\n",
      "Ep: 26, loss: 0.22, norm_all:1414.83, grp_sparsity: 0.70, acc1: 0.8381, norm_import: 1414.83, norm_redund: 0.00, num_grp_import: 865, num_grp_redund: 2015\n",
      "Ep: 27, loss: 0.21, norm_all:1419.18, grp_sparsity: 0.70, acc1: 0.8798, norm_import: 1419.18, norm_redund: 0.00, num_grp_import: 865, num_grp_redund: 2015\n",
      "Ep: 28, loss: 0.20, norm_all:1423.52, grp_sparsity: 0.70, acc1: 0.8103, norm_import: 1423.52, norm_redund: 0.00, num_grp_import: 865, num_grp_redund: 2015\n",
      "Ep: 29, loss: 0.19, norm_all:1427.37, grp_sparsity: 0.70, acc1: 0.8527, norm_import: 1427.37, norm_redund: 0.00, num_grp_import: 865, num_grp_redund: 2015\n",
      "Ep: 30, loss: 0.19, norm_all:1430.80, grp_sparsity: 0.70, acc1: 0.8972, norm_import: 1430.80, norm_redund: 0.00, num_grp_import: 865, num_grp_redund: 2015\n",
      "...",
      "Ep: 80, loss: 0.02, norm_all:1432.82, grp_sparsity: 0.70, acc1: 0.9274, norm_import: 1432.82, norm_redund: 0.00, num_grp_import: 865, num_grp_redund: 2015\n",
      "Ep: 81, loss: 0.01, norm_all:1431.79, grp_sparsity: 0.70, acc1: 0.9261, norm_import: 1431.79, norm_redund: 0.00, num_grp_import: 865, num_grp_redund: 2015\n",
      "Ep: 82, loss: 0.02, norm_all:1430.77, grp_sparsity: 0.70, acc1: 0.9275, norm_import: 1430.77, norm_redund: 0.00, num_grp_import: 865, num_grp_redund: 2015\n",
      "Ep: 83, loss: 0.01, norm_all:1429.75, grp_sparsity: 0.70, acc1: 0.9262, norm_import: 1429.75, norm_redund: 0.00, num_grp_import: 865, num_grp_redund: 2015\n",
      "Ep: 84, loss: 0.01, norm_all:1428.72, grp_sparsity: 0.70, acc1: 0.9274, norm_import: 1428.72, norm_redund: 0.00, num_grp_import: 865, num_grp_redund: 2015\n",
      "Ep: 85, loss: 0.01, norm_all:1427.71, grp_sparsity: 0.70, acc1: 0.9261, norm_import: 1427.71, norm_redund: 0.00, num_grp_import: 865, num_grp_redund: 2015\n",
      "Ep: 86, loss: 0.01, norm_all:1426.68, grp_sparsity: 0.70, acc1: 0.9255, norm_import: 1426.68, norm_redund: 0.00, num_grp_import: 865, num_grp_redund: 2015\n",
      "Ep: 87, loss: 0.01, norm_all:1425.65, grp_sparsity: 0.70, acc1: 0.9271, norm_import: 1425.65, norm_redund: 0.00, num_grp_import: 865, num_grp_redund: 2015\n",
      "Ep: 88, loss: 0.01, norm_all:1424.63, grp_sparsity: 0.70, acc1: 0.9258, norm_import: 1424.63, norm_redund: 0.00, num_grp_import: 865, num_grp_redund: 2015\n",
      "Ep: 89, loss: 0.01, norm_all:1423.61, grp_sparsity: 0.70, acc1: 0.9246, norm_import: 1423.61, norm_redund: 0.00, num_grp_import: 865, num_grp_redund: 2015\n",
      "Ep: 90, loss: 0.01, norm_all:1422.59, grp_sparsity: 0.70, acc1: 0.9273, norm_import: 1422.59, norm_redund: 0.00, num_grp_import: 865, num_grp_redund: 2015\n",
      "Ep: 91, loss: 0.01, norm_all:1421.57, grp_sparsity: 0.70, acc1: 0.9257, norm_import: 1421.57, norm_redund: 0.00, num_grp_import: 865, num_grp_redund: 2015\n",
      "Ep: 92, loss: 0.01, norm_all:1420.55, grp_sparsity: 0.70, acc1: 0.9275, norm_import: 1420.55, norm_redund: 0.00, num_grp_import: 865, num_grp_redund: 2015\n",
      "Ep: 93, loss: 0.01, norm_all:1419.52, grp_sparsity: 0.70, acc1: 0.9258, norm_import: 1419.52, norm_redund: 0.00, num_grp_import: 865, num_grp_redund: 2015\n",
      "Ep: 94, loss: 0.01, norm_all:1418.50, grp_sparsity: 0.70, acc1: 0.9268, norm_import: 1418.50, norm_redund: 0.00, num_grp_import: 865, num_grp_redund: 2015\n",
      "Ep: 95, loss: 0.01, norm_all:1417.47, grp_sparsity: 0.70, acc1: 0.9257, norm_import: 1417.47, norm_redund: 0.00, num_grp_import: 865, num_grp_redund: 2015\n",
      "Ep: 96, loss: 0.01, norm_all:1416.45, grp_sparsity: 0.70, acc1: 0.9267, norm_import: 1416.45, norm_redund: 0.00, num_grp_import: 865, num_grp_redund: 2015\n",
      "Ep: 97, loss: 0.01, norm_all:1415.43, grp_sparsity: 0.70, acc1: 0.9286, norm_import: 1415.43, norm_redund: 0.00, num_grp_import: 865, num_grp_redund: 2015\n",
      "Ep: 98, loss: 0.01, norm_all:1414.41, grp_sparsity: 0.70, acc1: 0.9291, norm_import: 1414.41, norm_redund: 0.00, num_grp_import: 865, num_grp_redund: 2015\n",
      "Ep: 99, loss: 0.01, norm_all:1414.30, grp_sparsity: 0.70, acc1: 0.9281, norm_import: 1414.30, norm_redund: 0.00, num_grp_import: 865, num_grp_redund: 2015\n"
     ]
    }
   ],
   "source": [
    "from utils.utils import check_accuracy\n",
    "\n",
    "max_epoch = 100\n",
    "model.cuda()\n",
    "criterion = torch.nn.CrossEntropyLoss()\n",
    "# Every 50 epochs, decay lr by 10.0\n",
    "lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=50, gamma=0.1) \n",
    "\n",
    "for epoch in range(max_epoch):\n",
    "    f_avg_val = 0.0\n",
    "    model.train()\n",
    "    lr_scheduler.step()\n",
    "    for X, y in trainloader:\n",
    "        X = X.cuda()\n",
    "        y = y.cuda()\n",
    "        y_pred = model.forward(X)\n",
    "        f = criterion(y_pred, y)\n",
    "        optimizer.zero_grad()\n",
    "        f.backward()\n",
    "        f_avg_val += f\n",
    "        optimizer.step()\n",
    "    group_sparsity, param_norm, _ = optimizer.compute_group_sparsity_param_norm()\n",
    "    norm_important, norm_redundant, num_grps_important, num_grps_redundant = optimizer.compute_norm_groups()\n",
    "    accuracy1, accuracy5 = check_accuracy(model, testloader)\n",
    "    f_avg_val = f_avg_val.cpu().item() / len(trainloader)\n",
    "    \n",
    "    print(\"Ep: {ep}, loss: {f:.2f}, norm_all:{param_norm:.2f}, grp_sparsity: {gs:.2f}, acc1: {acc1:.4f}, norm_import: {norm_import:.2f}, norm_redund: {norm_redund:.2f}, num_grp_import: {num_grps_import}, num_grp_redund: {num_grps_redund}\"\\\n",
    "         .format(ep=epoch, f=f_avg_val, param_norm=param_norm, gs=group_sparsity, acc1=accuracy1,\\\n",
    "         norm_import=norm_important, norm_redund=norm_redundant, num_grps_import=num_grps_important, num_grps_redund=num_grps_redundant\n",
    "        ))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb8dabf9-04cf-4d55-89e4-a17c96d17cec",
   "metadata": {},
   "source": [
    "### Step 5. Get compressed model in torch format"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "f46e6b6e-6286-4f57-816f-f94f6ce5c40b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# By default OTO will construct subnet by the last checkpoint. If intermedia ckpt reaches the best performance,\n",
    "# need to reinitialize OTO instance\n",
    "# oto = OTO(torch.load(ckpt_path), dummy_input)\n",
    "# then construct subnetwork\n",
    "oto.construct_subnet(out_dir='./cache')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da9cd3b6-1db9-473d-9585-27e08c669265",
   "metadata": {},
   "source": [
    "### (Optional) Check the compressed model size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "dc731c99-3155-472e-9462-65ed8c9bf4d7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Size of full model     :  0.041715689934790134 GBs\n",
      "Size of compress model :  0.0033666836097836494 GBs\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "\n",
    "full_model_size = os.stat(oto.full_group_sparse_model_path)\n",
    "compressed_model_size = os.stat(oto.compressed_model_path)\n",
    "print(\"Size of full model     : \", full_model_size.st_size / (1024 ** 3), \"GBs\")\n",
    "print(\"Size of compress model : \", compressed_model_size.st_size / (1024 ** 3), \"GBs\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c6d825e-46c8-49e9-b9f0-6f4eea4d8e2d",
   "metadata": {},
   "source": [
    "### (Optional) Check the compressed model accuracy\r\n",
    "#### # Both full and compressed model should return the exact same accuracy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "be56f693-3566-4b91-b5a7-39f327e61450",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Full model: Acc 1: 0.928, Acc 5: 0.998\n",
      "Compressed model: Acc 1: 0.928, Acc 5: 0.998\n"
     ]
    }
   ],
   "source": [
    "full_model = torch.load(oto.full_group_sparse_model_path)\n",
    "compressed_model = torch.load(oto.compressed_model_path)\n",
    "\n",
    "acc1_full, acc5_full = check_accuracy(full_model, testloader)\n",
    "print(\"Full model: Acc 1: {acc1}, Acc 5: {acc5}\".format(acc1=acc1_full, acc5=acc5_full))\n",
    "\n",
    "acc1_compressed, acc5_compressed = check_accuracy(compressed_model, testloader)\n",
    "print(\"Compressed model: Acc 1: {acc1}, Acc 5: {acc5}\".format(acc1=acc1_compressed, acc5=acc5_compressed))"
   ]
  }
 ],
 "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.11.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}