transform_example.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torchvision\n",
    "import torchvision.models as models\n",
    "\n",
    "from transformers.torchTransformer import TorchTransformer\n",
    "from transformers.quantize import QConv2d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = models.__dict__[\"resnet18\"]()\n",
    "model.cuda()\n",
    "model = model.eval()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Register layer to be transform"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "register <class 'torch.nn.modules.conv.Conv2d'> <class 'transformers.quantize.QConv2d'>\n"
     ]
    }
   ],
   "source": [
    "transformer = TorchTransformer()\n",
    "transformer.register(nn.Conv2d, QConv2d)\n",
    "model = transformer.trans_layers(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ResNet(\n",
       "  (conv1): QConv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)\n",
       "  (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (relu): ReLU(inplace=True)\n",
       "  (maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)\n",
       "  (layer1): Sequential(\n",
       "    (0): BasicBlock(\n",
       "      (conv1): QConv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "      (conv2): QConv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "    (1): BasicBlock(\n",
       "      (conv1): QConv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "      (conv2): QConv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "  )\n",
       "  (layer2): Sequential(\n",
       "    (0): BasicBlock(\n",
       "      (conv1): QConv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "      (conv2): QConv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (downsample): Sequential(\n",
       "        (0): QConv2d(64, 128, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "        (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (1): BasicBlock(\n",
       "      (conv1): QConv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "      (conv2): QConv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "  )\n",
       "  (layer3): Sequential(\n",
       "    (0): BasicBlock(\n",
       "      (conv1): QConv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "      (conv2): QConv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (downsample): Sequential(\n",
       "        (0): QConv2d(128, 256, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "        (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (1): BasicBlock(\n",
       "      (conv1): QConv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "      (conv2): QConv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "  )\n",
       "  (layer4): Sequential(\n",
       "    (0): BasicBlock(\n",
       "      (conv1): QConv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "      (conv2): QConv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (downsample): Sequential(\n",
       "        (0): QConv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "        (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (1): BasicBlock(\n",
       "      (conv1): QConv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "      (conv2): QConv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "  )\n",
       "  (avgpool): AdaptiveAvgPool2d(output_size=(1, 1))\n",
       "  (fc): Linear(in_features=512, out_features=1000, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "##########################################################################################\n",
      "Index| Layer (type)    | Bottoms         Output Shape              Param #        \n",
      "---------------------------------------------------------------------------\n",
      "    0| Data            |                 [(1, 3, 224, 224)]        0              \n",
      "---------------------------------------------------------------------------\n",
      "    1| QConv2d_1       | Data            [(1, 64, 112, 112)]       9408           \n",
      "---------------------------------------------------------------------------\n",
      "    2| BatchNorm2d_2   | QConv2d_1       [(1, 64, 112, 112)]       64             \n",
      "---------------------------------------------------------------------------\n",
      "    3| ReLU_3          | BatchNorm2d_2   [(1, 64, 112, 112)]       0              \n",
      "---------------------------------------------------------------------------\n",
      "    4| MaxPool2d_4     | ReLU_3          [(1, 64, 56, 56)]         0              \n",
      "---------------------------------------------------------------------------\n",
      "    5| QConv2d_5       | MaxPool2d_4     [(1, 64, 56, 56)]         36864          \n",
      "---------------------------------------------------------------------------\n",
      "    6| BatchNorm2d_6   | QConv2d_5       [(1, 64, 56, 56)]         64             \n",
      "---------------------------------------------------------------------------\n",
      "    7| ReLU_7          | BatchNorm2d_6   [(1, 64, 56, 56)]         0              \n",
      "---------------------------------------------------------------------------\n",
      "    8| QConv2d_8       | ReLU_7          [(1, 64, 56, 56)]         36864          \n",
      "---------------------------------------------------------------------------\n",
      "    9| BatchNorm2d_9   | QConv2d_8       [(1, 64, 56, 56)]         64             \n",
      "---------------------------------------------------------------------------\n",
      "   10| iadd_10         | BatchNorm2d_9   [(1, 64, 56, 56)]         0              \n",
      "     |                 | MaxPool2d_4                                              \n",
      "---------------------------------------------------------------------------\n",
      "   11| ReLU_11         | iadd_10         [(1, 64, 56, 56)]         0              \n",
      "---------------------------------------------------------------------------\n",
      "   12| QConv2d_12      | ReLU_11         [(1, 64, 56, 56)]         36864          \n",
      "---------------------------------------------------------------------------\n",
      "   13| BatchNorm2d_13  | QConv2d_12      [(1, 64, 56, 56)]         64             \n",
      "---------------------------------------------------------------------------\n",
      "   14| ReLU_14         | BatchNorm2d_13  [(1, 64, 56, 56)]         0              \n",
      "---------------------------------------------------------------------------\n",
      "   15| QConv2d_15      | ReLU_14         [(1, 64, 56, 56)]         36864          \n",
      "---------------------------------------------------------------------------\n",
      "   16| BatchNorm2d_16  | QConv2d_15      [(1, 64, 56, 56)]         64             \n",
      "---------------------------------------------------------------------------\n",
      "   17| iadd_17         | BatchNorm2d_16  [(1, 64, 56, 56)]         0              \n",
      "     |                 | ReLU_11                                                  \n",
      "---------------------------------------------------------------------------\n",
      "   18| ReLU_18         | iadd_17         [(1, 64, 56, 56)]         0              \n",
      "---------------------------------------------------------------------------\n",
      "   19| QConv2d_19      | ReLU_18         [(1, 128, 28, 28)]        73728          \n",
      "---------------------------------------------------------------------------\n",
      "   20| BatchNorm2d_20  | QConv2d_19      [(1, 128, 28, 28)]        128            \n",
      "---------------------------------------------------------------------------\n",
      "   21| ReLU_21         | BatchNorm2d_20  [(1, 128, 28, 28)]        0              \n",
      "---------------------------------------------------------------------------\n",
      "   22| QConv2d_22      | ReLU_21         [(1, 128, 28, 28)]        147456         \n",
      "---------------------------------------------------------------------------\n",
      "   23| BatchNorm2d_23  | QConv2d_22      [(1, 128, 28, 28)]        128            \n",
      "---------------------------------------------------------------------------\n",
      "   24| QConv2d_24      | ReLU_18         [(1, 128, 28, 28)]        8192           \n",
      "---------------------------------------------------------------------------\n",
      "   25| BatchNorm2d_25  | QConv2d_24      [(1, 128, 28, 28)]        128            \n",
      "---------------------------------------------------------------------------\n",
      "   26| iadd_26         | BatchNorm2d_23  [(1, 128, 28, 28)]        0              \n",
      "     |                 | BatchNorm2d_25                                           \n",
      "---------------------------------------------------------------------------\n",
      "   27| ReLU_27         | iadd_26         [(1, 128, 28, 28)]        0              \n",
      "---------------------------------------------------------------------------\n",
      "   28| QConv2d_28      | ReLU_27         [(1, 128, 28, 28)]        147456         \n",
      "---------------------------------------------------------------------------\n",
      "   29| BatchNorm2d_29  | QConv2d_28      [(1, 128, 28, 28)]        128            \n",
      "---------------------------------------------------------------------------\n",
      "   30| ReLU_30         | BatchNorm2d_29  [(1, 128, 28, 28)]        0              \n",
      "---------------------------------------------------------------------------\n",
      "   31| QConv2d_31      | ReLU_30         [(1, 128, 28, 28)]        147456         \n",
      "---------------------------------------------------------------------------\n",
      "   32| BatchNorm2d_32  | QConv2d_31      [(1, 128, 28, 28)]        128            \n",
      "---------------------------------------------------------------------------\n",
      "   33| iadd_33         | BatchNorm2d_32  [(1, 128, 28, 28)]        0              \n",
      "     |                 | ReLU_27                                                  \n",
      "---------------------------------------------------------------------------\n",
      "   34| ReLU_34         | iadd_33         [(1, 128, 28, 28)]        0              \n",
      "---------------------------------------------------------------------------\n",
      "   35| QConv2d_35      | ReLU_34         [(1, 256, 14, 14)]        294912         \n",
      "---------------------------------------------------------------------------\n",
      "   36| BatchNorm2d_36  | QConv2d_35      [(1, 256, 14, 14)]        256            \n",
      "---------------------------------------------------------------------------\n",
      "   37| ReLU_37         | BatchNorm2d_36  [(1, 256, 14, 14)]        0              \n",
      "---------------------------------------------------------------------------\n",
      "   38| QConv2d_38      | ReLU_37         [(1, 256, 14, 14)]        589824         \n",
      "---------------------------------------------------------------------------\n",
      "   39| BatchNorm2d_39  | QConv2d_38      [(1, 256, 14, 14)]        256            \n",
      "---------------------------------------------------------------------------\n",
      "   40| QConv2d_40      | ReLU_34         [(1, 256, 14, 14)]        32768          \n",
      "---------------------------------------------------------------------------\n",
      "   41| BatchNorm2d_41  | QConv2d_40      [(1, 256, 14, 14)]        256            \n",
      "---------------------------------------------------------------------------\n",
      "   42| iadd_42         | BatchNorm2d_39  [(1, 256, 14, 14)]        0              \n",
      "     |                 | BatchNorm2d_41                                           \n",
      "---------------------------------------------------------------------------\n",
      "   43| ReLU_43         | iadd_42         [(1, 256, 14, 14)]        0              \n",
      "---------------------------------------------------------------------------\n",
      "   44| QConv2d_44      | ReLU_43         [(1, 256, 14, 14)]        589824         \n",
      "---------------------------------------------------------------------------\n",
      "   45| BatchNorm2d_45  | QConv2d_44      [(1, 256, 14, 14)]        256            \n",
      "---------------------------------------------------------------------------\n",
      "   46| ReLU_46         | BatchNorm2d_45  [(1, 256, 14, 14)]        0              \n",
      "---------------------------------------------------------------------------\n",
      "   47| QConv2d_47      | ReLU_46         [(1, 256, 14, 14)]        589824         \n",
      "---------------------------------------------------------------------------\n",
      "   48| BatchNorm2d_48  | QConv2d_47      [(1, 256, 14, 14)]        256            \n",
      "---------------------------------------------------------------------------\n",
      "   49| iadd_49         | BatchNorm2d_48  [(1, 256, 14, 14)]        0              \n",
      "     |                 | ReLU_43                                                  \n",
      "---------------------------------------------------------------------------\n",
      "   50| ReLU_50         | iadd_49         [(1, 256, 14, 14)]        0              \n",
      "---------------------------------------------------------------------------\n",
      "   51| QConv2d_51      | ReLU_50         [(1, 512, 7, 7)]          1179648        \n",
      "---------------------------------------------------------------------------\n",
      "   52| BatchNorm2d_52  | QConv2d_51      [(1, 512, 7, 7)]          512            \n",
      "---------------------------------------------------------------------------\n",
      "   53| ReLU_53         | BatchNorm2d_52  [(1, 512, 7, 7)]          0              \n",
      "---------------------------------------------------------------------------\n",
      "   54| QConv2d_54      | ReLU_53         [(1, 512, 7, 7)]          2359296        \n",
      "---------------------------------------------------------------------------\n",
      "   55| BatchNorm2d_55  | QConv2d_54      [(1, 512, 7, 7)]          512            \n",
      "---------------------------------------------------------------------------\n",
      "   56| QConv2d_56      | ReLU_50         [(1, 512, 7, 7)]          131072         \n",
      "---------------------------------------------------------------------------\n",
      "   57| BatchNorm2d_57  | QConv2d_56      [(1, 512, 7, 7)]          512            \n",
      "---------------------------------------------------------------------------\n",
      "   58| iadd_58         | BatchNorm2d_55  [(1, 512, 7, 7)]          0              \n",
      "     |                 | BatchNorm2d_57                                           \n",
      "---------------------------------------------------------------------------\n",
      "   59| ReLU_59         | iadd_58         [(1, 512, 7, 7)]          0              \n",
      "---------------------------------------------------------------------------\n",
      "   60| QConv2d_60      | ReLU_59         [(1, 512, 7, 7)]          2359296        \n",
      "---------------------------------------------------------------------------\n",
      "   61| BatchNorm2d_61  | QConv2d_60      [(1, 512, 7, 7)]          512            \n",
      "---------------------------------------------------------------------------\n",
      "   62| ReLU_62         | BatchNorm2d_61  [(1, 512, 7, 7)]          0              \n",
      "---------------------------------------------------------------------------\n",
      "   63| QConv2d_63      | ReLU_62         [(1, 512, 7, 7)]          2359296        \n",
      "---------------------------------------------------------------------------\n",
      "   64| BatchNorm2d_64  | QConv2d_63      [(1, 512, 7, 7)]          512            \n",
      "---------------------------------------------------------------------------\n",
      "   65| iadd_65         | BatchNorm2d_64  [(1, 512, 7, 7)]          0              \n",
      "     |                 | ReLU_59                                                  \n",
      "---------------------------------------------------------------------------\n",
      "   66| ReLU_66         | iadd_65         [(1, 512, 7, 7)]          0              \n",
      "---------------------------------------------------------------------------\n",
      "   67| AdaptiveAvgPool2d_67 | ReLU_66         [(1, 512, 1, 1)]          0              \n",
      "---------------------------------------------------------------------------\n",
      "   68| torch.flatten_68 | AdaptiveAvgPool2d_67 [(1, 512)]                0              \n",
      "---------------------------------------------------------------------------\n",
      "   69| Linear_69       | torch.flatten_68 [(1, 1000)]               512000         \n",
      "---------------------------------------------------------------------------\n",
      "==================================================================================\n",
      "Total Trainable params: 11683712 \n",
      "Total Non-Trainable params: 0 \n",
      "Total params: 11683712 \n"
     ]
    }
   ],
   "source": [
    "input_tensor = torch.randn([1, 3, 224, 224]).cuda()\n",
    "model = model.cuda()\n",
    "transformer.summary(model, input_tensor = input_tensor)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}