convert_repvgg.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "from torchinfo import summary\n",
    "from backbones import get_model"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Init and load training model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = get_model('repvgg', dropout=0.0, fp16=False, num_features=512)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load RepVGG training checkpoint\n",
    "weight = torch.load('work_dirs/ms1mv3_repvgg/model.pt', map_location=torch.device('cpu'))\n",
    "model.load_state_dict(weight, strict=True)\n",
    "model = model.eval()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Stats"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.rand((1, 3, 112, 112))\n",
    "with torch.no_grad():\n",
    "    y_training = model(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "==========================================================================================\n",
       "Layer (type:depth-idx)                   Param #                   Mult-Adds\n",
       "==========================================================================================\n",
       "RepVGG                                   --                        --\n",
       "├─RepVGGBlock: 1-1                       --                        --\n",
       "│    └─Sequential: 2-1                   1,392                     4,064,352\n",
       "│    └─Sequential: 2-2                   240                       451,680\n",
       "│    └─Identity: 2-3                     --                        --\n",
       "│    └─ReLU: 2-4                         --                        --\n",
       "├─ModuleList: 1-2                        --                        --\n",
       "│    └─RepVGGBlock: 2-5                  23,232                    18,063,552\n",
       "│    └─RepVGGBlock: 2-6                  23,328                    18,063,648\n",
       "├─ModuleList: 1-3                        --                        --\n",
       "│    └─RepVGGBlock: 2-7                  46,464                    9,032,064\n",
       "│    └─RepVGGBlock: 2-8                  92,736                    18,063,936\n",
       "│    └─RepVGGBlock: 2-9                  92,736                    18,063,936\n",
       "│    └─RepVGGBlock: 2-10                 92,736                    18,063,936\n",
       "├─ModuleList: 1-4                        --                        --\n",
       "│    └─RepVGGBlock: 2-11                 185,088                   9,032,448\n",
       "│    └─RepVGGBlock: 2-12                 369,792                   18,064,512\n",
       "│    └─RepVGGBlock: 2-13                 369,792                   18,064,512\n",
       "│    └─RepVGGBlock: 2-14                 369,792                   18,064,512\n",
       "│    └─RepVGGBlock: 2-15                 369,792                   18,064,512\n",
       "│    └─RepVGGBlock: 2-16                 369,792                   18,064,512\n",
       "│    └─RepVGGBlock: 2-17                 369,792                   18,064,512\n",
       "│    └─RepVGGBlock: 2-18                 369,792                   18,064,512\n",
       "│    └─RepVGGBlock: 2-19                 369,792                   18,064,512\n",
       "│    └─RepVGGBlock: 2-20                 369,792                   18,064,512\n",
       "│    └─RepVGGBlock: 2-21                 369,792                   18,064,512\n",
       "│    └─RepVGGBlock: 2-22                 369,792                   18,064,512\n",
       "│    └─RepVGGBlock: 2-23                 369,792                   18,064,512\n",
       "│    └─RepVGGBlock: 2-24                 369,792                   18,064,512\n",
       "├─ModuleList: 1-5                        --                        --\n",
       "│    └─RepVGGBlock: 2-25                 2,462,720                 39,326,720\n",
       "├─AdaptiveAvgPool2d: 1-6                 --                        --\n",
       "├─Linear: 1-7                            655,360                   655,360\n",
       "├─BatchNorm1d: 1-8                       1,024                     1,024\n",
       "==========================================================================================\n",
       "Total params: 8,484,352\n",
       "Trainable params: 8,484,352\n",
       "Non-trainable params: 0\n",
       "Total mult-adds (M): 387.72\n",
       "==========================================================================================\n",
       "Input size (MB): 0.15\n",
       "Forward/backward pass size (MB): 16.24\n",
       "Params size (MB): 33.94\n",
       "Estimated Total Size (MB): 50.33\n",
       "=========================================================================================="
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "summary(model, input_size=(1, 3, 112, 112), col_names=[\"num_params\", \"mult_adds\"], depth=2)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Switch RepVGG to deployment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "for m in model.modules():\n",
    "    if hasattr(m, 'switch_to_deploy'):\n",
    "        m.switch_to_deploy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(model.state_dict(), 'work_dirs/ms1mv3_repvgg/model_deploy.pt')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Stats"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "==========================================================================================\n",
       "Layer (type:depth-idx)                   Param #                   Mult-Adds\n",
       "==========================================================================================\n",
       "RepVGG                                   --                        --\n",
       "├─RepVGGBlock: 1-1                       --                        --\n",
       "│    └─Conv2d: 2-1                       1,344                     4,214,784\n",
       "│    └─Identity: 2-2                     --                        --\n",
       "│    └─ReLU: 2-3                         --                        --\n",
       "├─ModuleList: 1-2                        --                        --\n",
       "│    └─RepVGGBlock: 2-4                  20,784                    16,294,656\n",
       "│    └─RepVGGBlock: 2-5                  20,784                    16,294,656\n",
       "├─ModuleList: 1-3                        --                        --\n",
       "│    └─RepVGGBlock: 2-6                  41,568                    8,147,328\n",
       "│    └─RepVGGBlock: 2-7                  83,040                    16,275,840\n",
       "│    └─RepVGGBlock: 2-8                  83,040                    16,275,840\n",
       "│    └─RepVGGBlock: 2-9                  83,040                    16,275,840\n",
       "├─ModuleList: 1-4                        --                        --\n",
       "│    └─RepVGGBlock: 2-10                 166,080                   8,137,920\n",
       "│    └─RepVGGBlock: 2-11                 331,968                   16,266,432\n",
       "│    └─RepVGGBlock: 2-12                 331,968                   16,266,432\n",
       "│    └─RepVGGBlock: 2-13                 331,968                   16,266,432\n",
       "│    └─RepVGGBlock: 2-14                 331,968                   16,266,432\n",
       "│    └─RepVGGBlock: 2-15                 331,968                   16,266,432\n",
       "│    └─RepVGGBlock: 2-16                 331,968                   16,266,432\n",
       "│    └─RepVGGBlock: 2-17                 331,968                   16,266,432\n",
       "│    └─RepVGGBlock: 2-18                 331,968                   16,266,432\n",
       "│    └─RepVGGBlock: 2-19                 331,968                   16,266,432\n",
       "│    └─RepVGGBlock: 2-20                 331,968                   16,266,432\n",
       "│    └─RepVGGBlock: 2-21                 331,968                   16,266,432\n",
       "│    └─RepVGGBlock: 2-22                 331,968                   16,266,432\n",
       "│    └─RepVGGBlock: 2-23                 331,968                   16,266,432\n",
       "├─ModuleList: 1-5                        --                        --\n",
       "│    └─RepVGGBlock: 2-24                 2,213,120                 35,409,920\n",
       "├─AdaptiveAvgPool2d: 1-6                 --                        --\n",
       "├─Linear: 1-7                            655,360                   655,360\n",
       "├─BatchNorm1d: 1-8                       1,024                     1,024\n",
       "==========================================================================================\n",
       "Total params: 7,684,768\n",
       "Trainable params: 7,684,768\n",
       "Non-trainable params: 0\n",
       "Total mult-adds (M): 349.45\n",
       "==========================================================================================\n",
       "Input size (MB): 0.15\n",
       "Forward/backward pass size (MB): 3.63\n",
       "Params size (MB): 30.74\n",
       "Estimated Total Size (MB): 34.52\n",
       "=========================================================================================="
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "summary(model, input_size=(1, 3, 112, 112), col_names=[\"num_params\", \"mult_adds\"], depth=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "with torch.no_grad():\n",
    "    y_deploy = model(x)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Check output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.allclose(y_training, y_deploy, rtol=1e-5, atol=1e-5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Euclidean distance: 3.739324683227352e-10\n"
     ]
    }
   ],
   "source": [
    "euclidena_dist = sum(((y_deploy - y_training)**2).reshape(512))\n",
    "print('Euclidean distance:', euclidena_dist.item())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cosine similarity: 1.0\n"
     ]
    }
   ],
   "source": [
    "cos = torch.nn.CosineSimilarity(dim=1)(y_training, y_deploy)\n",
    "print('Cosine similarity:', cos.item())"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MBF's stats"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_mbf = get_model('mbf', dropout=0.0, fp16=False, num_features=512)\n",
    "model_mbf = model_mbf.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "==============================================================================================================\n",
       "Layer (type:depth-idx)                                       Param #                   Mult-Adds\n",
       "==============================================================================================================\n",
       "MobileFaceNet                                                --                        --\n",
       "├─ModuleList: 1-1                                            --                        --\n",
       "│    └─ConvBlock: 2-1                                        3,840                     10,838,400\n",
       "│    └─ConvBlock: 2-2                                        2,688                     7,225,728\n",
       "│    └─DepthWise: 2-3                                        34,944                    65,129,472\n",
       "│    └─Residual: 2-4                                         139,776                   106,377,216\n",
       "│    └─DepthWise: 2-5                                        102,656                   38,988,800\n",
       "│    └─Residual: 2-6                                         812,544                   156,862,464\n",
       "│    └─DepthWise: 2-7                                        270,336                   32,342,016\n",
       "│    └─Residual: 2-8                                         270,848                   13,074,944\n",
       "├─ConvBlock: 1-2                                             --                        --\n",
       "│    └─Sequential: 2-9                                       132,608                   6,424,064\n",
       "├─GDC: 1-3                                                   --                        --\n",
       "│    └─Sequential: 2-10                                      289,280                   289,280\n",
       "==============================================================================================================\n",
       "Total params: 2,059,520\n",
       "Trainable params: 2,059,520\n",
       "Non-trainable params: 0\n",
       "Total mult-adds (M): 437.55\n",
       "==============================================================================================================\n",
       "Input size (MB): 0.15\n",
       "Forward/backward pass size (MB): 90.13\n",
       "Params size (MB): 8.24\n",
       "Estimated Total Size (MB): 98.52\n",
       "=============================================================================================================="
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "summary(model_mbf, input_size=(1, 3, 112, 112), col_names=[\"num_params\", \"mult_adds\"], depth=2)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "deep-learning",
   "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.8"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "c0a826c7eab53f154d9d3238481ec91124c15c94ad3afd1eabeb0e2f8a517073"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}