quickstart.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Profile a custom PyTorch model\n",
    "\n",
    "- Table of contents\n",
    "  - [Model-Level Profiling](#model-level-profiling)\n",
    "  - [Layer-Level Profiling](#layer-level-profiling)\n",
    "  - [Operator-Level Profiling](#operator-level-profiling)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from rich.table import Table\n",
    "from rich.console import Console\n",
    "\n",
    "console = Console()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# custom pytorch model\n",
    "class MyModel(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(MyModel, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(3, 3, kernel_size=3, padding=1)\n",
    "        self.conv2 = nn.Conv2d(3, 3, kernel_size=3, padding=1)\n",
    "        self.conv3 = nn.Conv2d(3, 3, kernel_size=3, padding=1)\n",
    "        self.conv4 = nn.Conv2d(3, 3, kernel_size=3, padding=1)\n",
    "\n",
    "    def forward(self, x1):\n",
    "        x1 = self.conv1(x1)\n",
    "        x1 = self.conv2(x1)\n",
    "        x1 = self.conv3(x1)\n",
    "        x1 = self.conv4(x1)\n",
    "        return x1\n",
    "\n",
    "# model = nn.Sequential(\n",
    "#     nn.Conv2d(3, 3, kernel_size=3, padding=1),\n",
    "#     nn.Conv2d(3, 3, kernel_size=3, padding=1),\n",
    "#     nn.Conv2d(3, 3, kernel_size=3, padding=1),\n",
    "#     nn.Conv2d(3, 3, kernel_size=3, padding=1)\n",
    "# ).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "def print_device(device):\n",
    "    \"\"\"print GPU info\"\"\"\n",
    "    table = Table(title=\"GPU Info\")\n",
    "    table.add_column(\"Device Name\", justify=\"left\", no_wrap=True)\n",
    "    table.add_column(\"Memory (GB)\", justify=\"left\", no_wrap=True)\n",
    "    table.add_column(\"SM Count\", justify=\"left\", no_wrap=True)\n",
    "    memory_gb = device.total_memory / 1024 / 1024 / 1024\n",
    "    table.add_row(device.name, '%.02f'%memory_gb, f\"{device.multi_processor_count}\")\n",
    "    console.print(table)\n",
    "    \n",
    "\n",
    "def print_model_profile(results):\n",
    "    \"\"\"print model-level profiling\"\"\"\n",
    "    # grid = Table.grid(expand=True)\n",
    "    table = Table(title=\"Model-Level Profiling\", show_lines=True)\n",
    "    table.add_column(\"Metrics\", justify=\"left\", no_wrap=True)\n",
    "    table.add_column(\"Cost\", justify=\"left\", no_wrap=True)\n",
    "    table.add_row(\"Latency (ms)\", f'{results[\"latency\"]}')\n",
    "    table.add_row(\"On-Device Inference (ms)\", f'{results[\"on_device_inference\"]}')\n",
    "    table.add_row(\"CPU-to-GPU Transfer (ms)\", f'{results[\"cpu_to_gpu_transfer\"]}')\n",
    "    table.add_row(\"GPU-to-CPU Transfer (ms)\", f'{results[\"gpu_to_cpu_transfer\"]}')\n",
    "    table.add_row(\"Maximum Memory (MB)\", f'{results[\"max_memory\"]}')\n",
    "    table.add_row(\"#Params\", f'{results[\"num_params\"]}')\n",
    "    table.add_row(\"#Macs\", f'{results[\"num_macs\"]}')\n",
    "    # table.add_column(\"Latency (ms)\", justify=\"left\", no_wrap=True)\n",
    "    # table.add_column(\"On-Device Inference (ms)\", justify=\"left\", no_wrap=True)\n",
    "    # table.add_column(\"CPU-to-GPU Transfer (ms)\", justify=\"left\", no_wrap=True)\n",
    "    # table.add_column(\"GPU-to-CPU Transfer (ms)\", justify=\"left\", no_wrap=True)\n",
    "    # table.add_column(\"Maximum Memory (MB)\", justify=\"left\", no_wrap=True)\n",
    "    # table.add_column(\"#Params\", justify=\"left\", no_wrap=True)\n",
    "    # table.add_column(\"#FLOPS\", justify=\"left\", no_wrap=True)\n",
    "    # # memory_gb = device.total_memory / 1024 / 1024 / 1024\n",
    "    # table.add_row(f'{results[\"latency\"]}', f'{results[\"on_device_inference\"]}', f'{results[\"cpu_to_gpu_transfer\"]}', f'{results[\"gpu_to_cpu_transfer\"]}', f'{results[\"max_memory\"]}', f'{results[\"num_params\"]}', f'{results[\"num_flops\"]}')\n",
    "    console.print(table)\n",
    "    \n",
    "    \n",
    "def print_layer_profile(results):\n",
    "    \"\"\"print layer-level profiling\"\"\"\n",
    "    # grid = Table.grid(expand=True)\n",
    "    table = Table(title=\"Layer-Level Profiling\", show_lines=True)\n",
    "    table.add_column(\"Layer\", justify=\"left\", no_wrap=True)\n",
    "    table.add_column(\"Latency (us)\", justify=\"left\", no_wrap=True)\n",
    "    table.add_column(\"#MACs (M)\", justify=\"left\", no_wrap=True)\n",
    "    for layer in results.keys():\n",
    "        layer_name = layer\n",
    "        latency = results[layer_name]['latency']\n",
    "        macs = results[layer_name]['macs']\n",
    "        table.add_row(layer_name, latency, macs)\n",
    "        # table.add_row(\"#MACs (M)\", f'{results[\"on_device_inference\"]}')\n",
    "    # table.add_row(\"CPU-to-GPU Transfer (ms)\", f'{results[\"cpu_to_gpu_transfer\"]}')\n",
    "    # table.add_row(\"GPU-to-CPU Transfer (ms)\", f'{results[\"gpu_to_cpu_transfer\"]}')\n",
    "    # table.add_row(\"Maximum Memory (MB)\", f'{results[\"max_memory\"]}')\n",
    "    # table.add_row(\"#Params\", f'{results[\"num_params\"]}')\n",
    "    # table.add_row(\"#Macs\", f'{results[\"num_macs\"]}')\n",
    "    # table.add_column(\"Latency (ms)\", justify=\"left\", no_wrap=True)\n",
    "    # table.add_column(\"On-Device Inference (ms)\", justify=\"left\", no_wrap=True)\n",
    "    # table.add_column(\"CPU-to-GPU Transfer (ms)\", justify=\"left\", no_wrap=True)\n",
    "    # table.add_column(\"GPU-to-CPU Transfer (ms)\", justify=\"left\", no_wrap=True)\n",
    "    # table.add_column(\"Maximum Memory (MB)\", justify=\"left\", no_wrap=True)\n",
    "    # table.add_column(\"#Params\", justify=\"left\", no_wrap=True)\n",
    "    # table.add_column(\"#FLOPS\", justify=\"left\", no_wrap=True)\n",
    "    # # memory_gb = device.total_memory / 1024 / 1024 / 1024\n",
    "    # table.add_row(f'{results[\"latency\"]}', f'{results[\"on_device_inference\"]}', f'{results[\"cpu_to_gpu_transfer\"]}', f'{results[\"gpu_to_cpu_transfer\"]}', f'{results[\"max_memory\"]}', f'{results[\"num_params\"]}', f'{results[\"num_flops\"]}')\n",
    "    console.print(table)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-style: italic\">                      GPU Info                      </span>\n",
       "┏━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━┳━━━━━━━━━━┓\n",
       "┃<span style=\"font-weight: bold\"> Device Name             </span>┃<span style=\"font-weight: bold\"> Memory (GB) </span>┃<span style=\"font-weight: bold\"> SM Count </span>┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━╇━━━━━━━━━━┩\n",
       "│ NVIDIA GeForce GTX 1070 │ 7.92        │ 16       │\n",
       "└─────────────────────────┴─────────────┴──────────┘\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[3m                      GPU Info                      \u001b[0m\n",
       "┏━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━┳━━━━━━━━━━┓\n",
       "┃\u001b[1m \u001b[0m\u001b[1mDevice Name            \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mMemory (GB)\u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mSM Count\u001b[0m\u001b[1m \u001b[0m┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━╇━━━━━━━━━━┩\n",
       "│ NVIDIA GeForce GTX 1070 │ 7.92        │ 16       │\n",
       "└─────────────────────────┴─────────────┴──────────┘\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "device_info = torch.cuda.get_device_properties(device)\n",
    "if device == \"cuda\":\n",
    "    print_device(device_info)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = MyModel().to(device)\n",
    "model.eval()\n",
    "sample = torch.randn(1, 3, 224, 224).to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model-Level Profiling\n",
    "Latency, #Params, #MAC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from omegaconf import OmegaConf\n",
    "from pytorch_benchmark import benchmark"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Warming up with batch_size=1: 100%|██████████| 1/1 [00:00<00:00,  9.24it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Warning! No positional inputs found for a module, assuming batch size is 1.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Warming up with batch_size=1: 100%|██████████| 10/10 [00:00<00:00, 829.37it/s]\n",
      "Measuring inference for batch_size=1: 100%|██████████| 100/100 [00:00<00:00, 880.38it/s]\n",
      "Unable to measure energy consumption. Device must be a NVIDIA Jetson.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'machine_info': {'system': {'system': 'Linux',\n",
       "   'node': 'jason-Alienware-17-R5',\n",
       "   'release': '5.15.0-107-generic'},\n",
       "  'cpu': {'model': 'Intel(R) Core(TM) i7-8750H CPU @ 2.20GHz',\n",
       "   'architecture': 'x86_64',\n",
       "   'cores': {'physical': 6, 'total': 12},\n",
       "   'frequency': '4.10 GHz'},\n",
       "  'memory': {'total': '31.20 GB', 'used': '16.20 GB', 'available': '13.66 GB'},\n",
       "  'gpus': [{'name': 'NVIDIA GeForce GTX 1070', 'memory': '8192.0 MB'}]},\n",
       " 'device': 'cuda',\n",
       " 'params': 336,\n",
       " 'flops': 16859136,\n",
       " 'timing': {'batch_size_1': {'on_device_inference': {'metrics': {'batches_per_second_mean': -1.7964813225007352,\n",
       "     'batches_per_second_std': 0.3060285807350442,\n",
       "     'batches_per_second_min': -2.1350003521808074,\n",
       "     'batches_per_second_max': -0.27640191645860046,\n",
       "     'seconds_per_batch_mean': -0.5996054401993751,\n",
       "     'seconds_per_batch_std': 0.32559099459601043,\n",
       "     'seconds_per_batch_min': -3.617919921875,\n",
       "     'seconds_per_batch_max': -0.4683839976787567},\n",
       "    'human_readable': {'batches_per_second': '-1.80 +/- 0.31 [-2.14, -0.28]',\n",
       "     'batch_latency': '-599605.440 us +/- 325.591 ms [-3617919.922 us, -468383.998 us]'}},\n",
       "   'cpu_to_gpu': {'metrics': {'batches_per_second_mean': 4021.6803650766733,\n",
       "     'batches_per_second_std': 749.477440074296,\n",
       "     'batches_per_second_min': 1913.4598540145985,\n",
       "     'batches_per_second_max': 5242.88,\n",
       "     'seconds_per_batch_mean': 0.00025937557220458987,\n",
       "     'seconds_per_batch_std': 6.020041739524279e-05,\n",
       "     'seconds_per_batch_min': 0.00019073486328125,\n",
       "     'seconds_per_batch_max': 0.000522613525390625},\n",
       "    'human_readable': {'batches_per_second': '4.02 K +/- 749.48 [1.91 K, 5.24 K]',\n",
       "     'batch_latency': '259.376 us +/- 60.200 us [190.735 us, 522.614 us]'}},\n",
       "   'gpu_to_cpu': {'metrics': {'batches_per_second_mean': 4161.507053397585,\n",
       "     'batches_per_second_std': 790.2972486382257,\n",
       "     'batches_per_second_min': 1807.1107281344248,\n",
       "     'batches_per_second_max': 5315.974651457541,\n",
       "     'seconds_per_batch_mean': 0.00025429964065551755,\n",
       "     'seconds_per_batch_std': 7.718047926361167e-05,\n",
       "     'seconds_per_batch_min': 0.0001881122589111328,\n",
       "     'seconds_per_batch_max': 0.0005533695220947266},\n",
       "    'human_readable': {'batches_per_second': '4.16 K +/- 790.30 [1.81 K, 5.32 K]',\n",
       "     'batch_latency': '254.300 us +/- 77.180 us [188.112 us, 553.370 us]'}},\n",
       "   'total': {'metrics': {'batches_per_second_mean': 929.951976321434,\n",
       "     'batches_per_second_std': 150.56069336041585,\n",
       "     'batches_per_second_min': 243.57166085946574,\n",
       "     'batches_per_second_max': 1119.0779082177162,\n",
       "     'seconds_per_batch_mean': 0.0011231470108032226,\n",
       "     'seconds_per_batch_std': 0.0003495611516817312,\n",
       "     'seconds_per_batch_min': 0.0008935928344726562,\n",
       "     'seconds_per_batch_max': 0.004105567932128906},\n",
       "    'human_readable': {'batches_per_second': '929.95 +/- 150.56 [243.57, 1.12 K]',\n",
       "     'batch_latency': '1.123 ms +/- 349.561 us [893.593 us, 4.106 ms]'}}}},\n",
       " 'memory': {'batch_size_1': {'pre_inference_bytes': 606208,\n",
       "   'max_inference_bytes': 2412544,\n",
       "   'post_inference_bytes': 606208,\n",
       "   'pre_inference': '592.00 KB',\n",
       "   'max_inference': '2.30 MB',\n",
       "   'post_inference': '592.00 KB'}}}"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "with torch.no_grad():\n",
    "    results = benchmark(model, sample, num_runs=100)\n",
    "\n",
    "\n",
    "results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-style: italic\">         Model-Level Profiling         </span>\n",
       "┏━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━┓\n",
       "┃<span style=\"font-weight: bold\"> Metrics                  </span>┃<span style=\"font-weight: bold\"> Cost     </span>┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━┩\n",
       "│ Latency (ms)             │ 0.991    │\n",
       "├──────────────────────────┼──────────┤\n",
       "│ On-Device Inference (ms) │ 0.554    │\n",
       "├──────────────────────────┼──────────┤\n",
       "│ CPU-to-GPU Transfer (ms) │ 0.227    │\n",
       "├──────────────────────────┼──────────┤\n",
       "│ GPU-to-CPU Transfer (ms) │ 0.21     │\n",
       "├──────────────────────────┼──────────┤\n",
       "│ Maximum Memory (MB)      │ 2.3      │\n",
       "├──────────────────────────┼──────────┤\n",
       "│ #Params                  │ 336      │\n",
       "├──────────────────────────┼──────────┤\n",
       "│ #Macs                    │ 16859136 │\n",
       "└──────────────────────────┴──────────┘\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[3m         Model-Level Profiling         \u001b[0m\n",
       "┏━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━┓\n",
       "┃\u001b[1m \u001b[0m\u001b[1mMetrics                 \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mCost    \u001b[0m\u001b[1m \u001b[0m┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━┩\n",
       "│ Latency (ms)             │ 0.991    │\n",
       "├──────────────────────────┼──────────┤\n",
       "│ On-Device Inference (ms) │ 0.554    │\n",
       "├──────────────────────────┼──────────┤\n",
       "│ CPU-to-GPU Transfer (ms) │ 0.227    │\n",
       "├──────────────────────────┼──────────┤\n",
       "│ GPU-to-CPU Transfer (ms) │ 0.21     │\n",
       "├──────────────────────────┼──────────┤\n",
       "│ Maximum Memory (MB)      │ 2.3      │\n",
       "├──────────────────────────┼──────────┤\n",
       "│ #Params                  │ 336      │\n",
       "├──────────────────────────┼──────────┤\n",
       "│ #Macs                    │ 16859136 │\n",
       "└──────────────────────────┴──────────┘\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# find metrics in print(results)\n",
    "# table.add_row(results[\"latency\"], results[\"on_device_inference\"], results[\"cpu_to_gpu_transfer\"], results[\"gpu_to_cpu_transfer\"], results[\"max_memory\"], results[\"num_params\"], results[\"num_flops\"])\n",
    "model_profile = {\n",
    "    'latency': 0.991,\n",
    "    'on_device_inference': 0.991 - 0.227 - 0.210,\n",
    "    'cpu_to_gpu_transfer': 0.227,\n",
    "    'gpu_to_cpu_transfer': 0.210,\n",
    "    'max_memory': 2.30,\n",
    "    'num_params': 336,\n",
    "    'num_macs': 16859136\n",
    "}\n",
    "\n",
    "print_model_profile(model_profile)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Layer-Level Profiling\n",
    "Latency, activations/memory, #MAC/FLOPS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pytorch_memlab import MemReporter\n",
    "from flops_profiler.profiler import get_model_profile"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Element type                                            Size  Used MEM\n",
      "-------------------------------------------------------------------------------\n",
      "Storage on cuda:0\n",
      "Tensor0                                     (1, 3, 224, 224)   588.00K\n",
      "Tensor1                                     (1, 3, 224, 224)   588.00K\n",
      "Tensor2                                     (1, 3, 224, 224)   588.00K\n",
      "conv1.weight                                    (3, 3, 3, 3)   512.00B\n",
      "conv1.bias                                              (3,)   512.00B\n",
      "conv2.weight                                    (3, 3, 3, 3)   512.00B\n",
      "conv2.bias                                              (3,)   512.00B\n",
      "conv3.weight                                    (3, 3, 3, 3)   512.00B\n",
      "conv3.bias                                              (3,)   512.00B\n",
      "conv4.weight                                    (3, 3, 3, 3)   512.00B\n",
      "conv4.bias                                              (3,)   512.00B\n",
      "-------------------------------------------------------------------------------\n",
      "Total Tensors: 451920 \tUsed Memory: 1.73M\n",
      "The allocated memory on cuda:0: 1.73M\n",
      "-------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/jason/miniconda3/envs/pytorch_profiler/lib/python3.9/site-packages/pytorch_memlab/mem_reporter.py:65: FutureWarning: `torch.distributed.reduce_op` is deprecated, please use `torch.distributed.ReduceOp` instead\n",
      "  tensors = [obj for obj in objects if isinstance(obj, torch.Tensor)]\n",
      "/home/jason/miniconda3/envs/pytorch_profiler/lib/python3.9/site-packages/pytorch_memlab/mem_reporter.py:95: UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly.  To access UntypedStorage directly, use tensor.untyped_storage() instead of tensor.storage()\n",
      "  fact_numel = tensor.storage().size()\n"
     ]
    }
   ],
   "source": [
    "# profile memory and #activations\n",
    "model = MyModel().to(device)\n",
    "model.eval()\n",
    "\n",
    "reporter = MemReporter(model)\n",
    "inp = torch.randn(1, 3, 224, 224).to(device)\n",
    "with torch.no_grad():\n",
    "    out = model(inp)\n",
    "    \n",
    "reporter.report(verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "-------------------------- Flops Profiler --------------------------\n",
      "Profile on Device: cuda:0\n",
      "Profile Summary at step 10:\n",
      "Notations:\n",
      "data parallel size (dp_size), model parallel size(mp_size),\n",
      "number of parameters (params), number of multiply-accumulate operations(MACs),\n",
      "number of floating-point operations (flops), floating-point operations per second (FLOPS),\n",
      "fwd latency (forward propagation latency), bwd latency (backward propagation latency),\n",
      "step (weights update latency), iter latency (sum of fwd, bwd and step latency)\n",
      "\n",
      "params per device:                                            336     \n",
      "params of model = params per device * mp_size:                336     \n",
      "fwd MACs per device:                                          16.26 MMACs\n",
      "fwd flops per device:                                         33.12 M \n",
      "fwd flops of model = fwd flops per device * mp_size:          33.12 M \n",
      "fwd latency:                                                  1.25 ms \n",
      "fwd FLOPS per device = fwd flops per device / fwd latency:    26.47 GFLOPS\n",
      "\n",
      "----------------------------- Aggregated Profile per Device -----------------------------\n",
      "Top 1 modules in terms of params, flops, MACs or duration at different model depths:\n",
      "depth 0:\n",
      "    params      - {'MyModel': '336'}\n",
      "    flops       - {'MyModel': '33.12 M'}\n",
      "    MACs        - {'MyModel': '16.26 MMACs'}\n",
      "    fwd latency - {'MyModel': '1.25 ms'}\n",
      "\n",
      "------------------------------ Detailed Profile per Device ------------------------------\n",
      "Each module profile is listed after its name in the following order: \n",
      "params, percentage of total params, MACs, percentage of total MACs, fwd latency, percentage of total fwd latency, fwd FLOPS\n",
      "\n",
      "Note: 1. A module can have torch.nn.module or torch.nn.functional to compute logits (e.g. CrossEntropyLoss). They are not counted as submodules, thus not to be printed out. However they make up the difference between a parent's MACs (or latency) and the sum of its submodules'.\n",
      "2. Number of floating-point operations is a theoretical estimation, thus FLOPS computed using that could be larger than the maximum system throughput.\n",
      "3. The fwd latency listed in the top module's profile is directly captured at the module forward function in PyTorch.\n",
      "\n",
      "MyModel(\n",
      "  module = {'param': '336', 'flops': '33.12 M', 'macs': '16.26 MMACs', 'duration': '1.25 ms', 'FLOPS': '26.47 GFLOPS', 'params%': '100.00%', 'flops%': '100.00%', 'macs%': '100.00%', 'duration%': '100.00%'}, functionals = {'conv2d': {'flops': '33.12 M', 'macs': '16.26 MMACs', 'duration': '605.79 us', 'FLOPS': '54.67 GFLOPS', 'flops%': '100.00%', 'macs%': '100.00%', 'duration%/allfuncs': '100.00%', 'duration%/e2e': '48.42%'}}, functionals_duration = 605.79 us, \n",
      "  (conv1): Conv2d(module = {'param': '84', 'flops': '8.28 M', 'macs': '4.06 MMACs', 'duration': '427.96 us', 'FLOPS': '19.35 GFLOPS', 'params%': '25.00%', 'flops%': '25.00%', 'macs%': '25.00%', 'duration%': '34.20%'}, functionals = {'conv2d': {'flops': '8.28 M', 'macs': '4.06 MMACs', 'duration': '209.92 us', 'FLOPS': '39.44 GFLOPS', 'flops%': '25.00%', 'macs%': '25.00%', 'duration%/allfuncs': '34.65%', 'duration%/e2e': '16.78%'}}, functionals_duration = 209.92 us, 3, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "  (conv2): Conv2d(module = {'param': '84', 'flops': '8.28 M', 'macs': '4.06 MMACs', 'duration': '275.37 us', 'FLOPS': '30.06 GFLOPS', 'params%': '25.00%', 'flops%': '25.00%', 'macs%': '25.00%', 'duration%': '22.01%'}, functionals = {'conv2d': {'flops': '8.28 M', 'macs': '4.06 MMACs', 'duration': '119.81 us', 'FLOPS': '69.1 GFLOPS', 'flops%': '25.00%', 'macs%': '25.00%', 'duration%/allfuncs': '19.78%', 'duration%/e2e': '9.58%'}}, functionals_duration = 119.81 us, 3, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "  (conv3): Conv2d(module = {'param': '84', 'flops': '8.28 M', 'macs': '4.06 MMACs', 'duration': '269.41 us', 'FLOPS': '30.73 GFLOPS', 'params%': '25.00%', 'flops%': '25.00%', 'macs%': '25.00%', 'duration%': '21.53%'}, functionals = {'conv2d': {'flops': '8.28 M', 'macs': '4.06 MMACs', 'duration': '156.26 us', 'FLOPS': '52.98 GFLOPS', 'flops%': '25.00%', 'macs%': '25.00%', 'duration%/allfuncs': '25.79%', 'duration%/e2e': '12.49%'}}, functionals_duration = 156.26 us, 3, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "  (conv4): Conv2d(module = {'param': '84', 'flops': '8.28 M', 'macs': '4.06 MMACs', 'duration': '278.47 us', 'FLOPS': '29.73 GFLOPS', 'params%': '25.00%', 'flops%': '25.00%', 'macs%': '25.00%', 'duration%': '22.26%'}, functionals = {'conv2d': {'flops': '8.28 M', 'macs': '4.06 MMACs', 'duration': '119.81 us', 'FLOPS': '69.1 GFLOPS', 'flops%': '25.00%', 'macs%': '25.00%', 'duration%/allfuncs': '19.78%', 'duration%/e2e': '9.58%'}}, functionals_duration = 119.81 us, 3, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      ")\n",
      "------------------------------------------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "# profile latency and #MACs\n",
    "with torch.no_grad():\n",
    "    model = MyModel().to(device)\n",
    "    model.eval()\n",
    "    batch_size = 1\n",
    "    flops, macs, params = get_model_profile(model=model, # model\n",
    "                                    input_shape=(batch_size, 3, 224, 224), # input shape to the model. If specified, the model takes a tensor with this shape as the only positional argument.\n",
    "                                    args=None, # list of positional arguments to the model.\n",
    "                                    kwargs=None, # dictionary of keyword arguments to the model.\n",
    "                                    print_profile=True, # prints the model graph with the measured profile attached to each module\n",
    "                                    detailed=True, # print the detailed profile\n",
    "                                    module_depth=-1, # depth into the nested modules, with -1 being the inner most modules\n",
    "                                    top_modules=1, # the number of top modules to print aggregated profile\n",
    "                                    warm_up=10, # the number of warm-ups before measuring the time of each module\n",
    "                                    as_string=True, # print raw numbers (e.g. 1000) or as human-readable strings (e.g. 1k)\n",
    "                                    output_file=None, # path to the output file. If None, the profiler prints to stdout.\n",
    "                                    ignore_modules=None, # the list of modules to ignore in the profiling\n",
    "                                    func_name='forward') # the function name to profile, \"forward\" by default, for huggingface generative models, `generate` is used"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-style: italic\">       Layer-Level Profiling        </span>\n",
       "┏━━━━━━━┳━━━━━━━━━━━━━━┳━━━━━━━━━━━┓\n",
       "┃<span style=\"font-weight: bold\"> Layer </span>┃<span style=\"font-weight: bold\"> Latency (us) </span>┃<span style=\"font-weight: bold\"> #MACs (M) </span>┃\n",
       "┡━━━━━━━╇━━━━━━━━━━━━━━╇━━━━━━━━━━━┩\n",
       "│ conv1 │ 427.96       │ 4.06      │\n",
       "├───────┼──────────────┼───────────┤\n",
       "│ conv2 │ 275.37       │ 4.06      │\n",
       "├───────┼──────────────┼───────────┤\n",
       "│ conv3 │ 269.41       │ 4.06      │\n",
       "├───────┼──────────────┼───────────┤\n",
       "│ conv4 │ 278.47       │ 4.06      │\n",
       "└───────┴──────────────┴───────────┘\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[3m       Layer-Level Profiling        \u001b[0m\n",
       "┏━━━━━━━┳━━━━━━━━━━━━━━┳━━━━━━━━━━━┓\n",
       "┃\u001b[1m \u001b[0m\u001b[1mLayer\u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mLatency (us)\u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1m#MACs (M)\u001b[0m\u001b[1m \u001b[0m┃\n",
       "┡━━━━━━━╇━━━━━━━━━━━━━━╇━━━━━━━━━━━┩\n",
       "│ conv1 │ 427.96       │ 4.06      │\n",
       "├───────┼──────────────┼───────────┤\n",
       "│ conv2 │ 275.37       │ 4.06      │\n",
       "├───────┼──────────────┼───────────┤\n",
       "│ conv3 │ 269.41       │ 4.06      │\n",
       "├───────┼──────────────┼───────────┤\n",
       "│ conv4 │ 278.47       │ 4.06      │\n",
       "└───────┴──────────────┴───────────┘\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "layer_profile = {\n",
    "    'conv1': {'macs': '4.06', 'latency': '427.96'},\n",
    "    'conv2': {'macs': '4.06', 'latency': '275.37'},\n",
    "    'conv3': {'macs': '4.06', 'latency': '269.41'},\n",
    "    'conv4': {'macs': '4.06', 'latency': '278.47'}\n",
    "}\n",
    "print_layer_profile(layer_profile)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Operator-Level Profiling\n",
    "Detailed execution graph and the most time-consuming operators"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.profiler import profile, record_function, ProfilerActivity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = MyModel().to(device)\n",
    "model.eval()\n",
    "sample_inp = torch.randn(1, 3, 224, 224).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], profile_memory=True, record_shapes=True) as prof:\n",
    "    model(sample_inp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  \n",
      "                                                   Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg     Self CUDA   Self CUDA %    CUDA total  CUDA time avg       CPU Mem  Self CPU Mem      CUDA Mem  Self CUDA Mem    # of Calls  \n",
      "-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  \n",
      "                                           aten::conv2d         0.43%      29.150us        95.01%       6.376ms       1.594ms       0.000us         0.00%     474.143us     118.536us           0 b           0 b       2.30 Mb           0 b             4  \n",
      "                                      aten::convolution         1.41%      94.647us        94.57%       6.346ms       1.587ms       0.000us         0.00%     474.143us     118.536us           0 b           0 b       2.30 Mb           0 b             4  \n",
      "                                     aten::_convolution         1.45%      97.074us        93.16%       6.252ms       1.563ms       0.000us         0.00%     474.143us     118.536us           0 b           0 b       2.30 Mb           0 b             4  \n",
      "                                aten::cudnn_convolution        36.63%       2.458ms        89.61%       6.014ms       1.503ms     445.471us        93.95%     445.471us     111.368us           0 b           0 b       2.30 Mb       2.30 Mb             4  \n",
      "void cudnn::cnn::conv2d_grouped_direct_kernel<false,...         0.00%       0.000us         0.00%       0.000us       0.000us     445.471us        93.95%     445.471us     111.368us           0 b           0 b           0 b           0 b             4  \n",
      "                                             aten::add_         1.01%      67.899us         1.77%     118.737us      29.684us      28.672us         6.05%      28.672us       7.168us           0 b           0 b           0 b           0 b             4  \n",
      "void at::native::elementwise_kernel<128, 2, at::nati...         0.00%       0.000us         0.00%       0.000us       0.000us      28.672us         6.05%      28.672us       7.168us           0 b           0 b           0 b           0 b             4  \n",
      "                                        cudaEventRecord         0.30%      19.967us         0.30%      19.967us       4.992us       0.000us         0.00%       0.000us       0.000us           0 b           0 b           0 b           0 b             4  \n",
      "                                  cudaStreamIsCapturing         0.08%       5.231us         0.08%       5.231us       1.308us       0.000us         0.00%       0.000us       0.000us           0 b           0 b           0 b           0 b             4  \n",
      "                                  cudaStreamGetPriority         0.07%       4.832us         0.07%       4.832us       1.208us       0.000us         0.00%       0.000us       0.000us           0 b           0 b           0 b           0 b             4  \n",
      "-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  \n",
      "Self CPU time total: 6.711ms\n",
      "Self CUDA time total: 474.143us\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(prof.key_averages().table(sort_by=\"cuda_time_total\", row_limit=10))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  \n",
      "                                                   Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg     Self CUDA   Self CUDA %    CUDA total  CUDA time avg       CPU Mem  Self CPU Mem      CUDA Mem  Self CUDA Mem    # of Calls  \n",
      "-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  \n",
      "                                aten::cudnn_convolution        36.63%       2.458ms        89.61%       6.014ms       1.503ms     445.471us        93.95%     445.471us     111.368us           0 b           0 b       2.30 Mb       2.30 Mb             4  \n",
      "                                           aten::conv2d         0.43%      29.150us        95.01%       6.376ms       1.594ms       0.000us         0.00%     474.143us     118.536us           0 b           0 b       2.30 Mb           0 b             4  \n",
      "                                      aten::convolution         1.41%      94.647us        94.57%       6.346ms       1.587ms       0.000us         0.00%     474.143us     118.536us           0 b           0 b       2.30 Mb           0 b             4  \n",
      "                                     aten::_convolution         1.45%      97.074us        93.16%       6.252ms       1.563ms       0.000us         0.00%     474.143us     118.536us           0 b           0 b       2.30 Mb           0 b             4  \n",
      "                                        cudaEventRecord         0.30%      19.967us         0.30%      19.967us       4.992us       0.000us         0.00%       0.000us       0.000us           0 b           0 b           0 b           0 b             4  \n",
      "                                  cudaStreamIsCapturing         0.08%       5.231us         0.08%       5.231us       1.308us       0.000us         0.00%       0.000us       0.000us           0 b           0 b           0 b           0 b             4  \n",
      "                                  cudaStreamGetPriority         0.07%       4.832us         0.07%       4.832us       1.208us       0.000us         0.00%       0.000us       0.000us           0 b           0 b           0 b           0 b             4  \n",
      "                       cudaDeviceGetStreamPriorityRange         0.05%       3.200us         0.05%       3.200us       0.800us       0.000us         0.00%       0.000us       0.000us           0 b           0 b           0 b           0 b             4  \n",
      "                                       cudaLaunchKernel        53.24%       3.573ms        53.24%       3.573ms     446.615us       0.000us         0.00%       0.000us       0.000us           0 b           0 b           0 b           0 b             8  \n",
      "                                          aten::reshape         0.17%      11.166us         0.33%      22.413us       5.603us       0.000us         0.00%       0.000us       0.000us           0 b           0 b           0 b           0 b             4  \n",
      "-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  \n",
      "Self CPU time total: 6.711ms\n",
      "Self CUDA time total: 474.143us\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(prof.key_averages().table(sort_by=\"self_cuda_memory_usage\", row_limit=10))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  --------------------------------------------------------------------------------  \n",
      "                                                   Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg     Self CUDA   Self CUDA %    CUDA total  CUDA time avg    # of Calls                                                                      Input Shapes  \n",
      "-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  --------------------------------------------------------------------------------  \n",
      "                                           aten::conv2d         0.23%      27.473us        98.87%      12.037ms       3.009ms       0.000us         0.00%     474.717us     118.679us             4                             [[1, 3, 224, 224], [3, 3, 3, 3], [3], [], [], [], []]  \n",
      "                                      aten::convolution         0.88%     106.680us        98.65%      12.010ms       3.002ms       0.000us         0.00%     474.717us     118.679us             4                     [[1, 3, 224, 224], [3, 3, 3, 3], [3], [], [], [], [], [], []]  \n",
      "                                     aten::_convolution         0.76%      92.662us        97.77%      11.903ms       2.976ms       0.000us         0.00%     474.717us     118.679us             4     [[1, 3, 224, 224], [3, 3, 3, 3], [3], [], [], [], [], [], [], [], [], [], []]  \n",
      "                                aten::cudnn_convolution        26.10%       3.178ms        95.73%      11.654ms       2.914ms     444.093us        93.55%     444.093us     111.023us             4                      [[1, 3, 224, 224], [3, 3, 3, 3], [], [], [], [], [], [], []]  \n",
      "                                       cudaLaunchKernel        69.82%       8.500ms        69.82%       8.500ms       1.062ms       0.000us         0.00%       0.000us       0.000us             8                                                                                []  \n",
      "                                  cudaDeviceSynchronize         1.13%     137.236us         1.13%     137.236us     137.236us       0.000us         0.00%       0.000us       0.000us             1                                                                                []  \n",
      "                                             aten::add_         0.60%      73.114us         1.08%     132.024us      33.006us      30.624us         6.45%      30.624us       7.656us             4                                              [[1, 3, 224, 224], [1, 3, 1, 1], []]  \n",
      "                                          aten::reshape         0.10%      11.700us         0.20%      24.108us       6.027us       0.000us         0.00%       0.000us       0.000us             4                                                                         [[3], []]  \n",
      "                                        cudaEventRecord         0.17%      21.054us         0.17%      21.054us       5.264us       0.000us         0.00%       0.000us       0.000us             4                                                                                []  \n",
      "                                             aten::view         0.10%      12.408us         0.10%      12.408us       3.102us       0.000us         0.00%       0.000us       0.000us             4                                                                         [[3], []]  \n",
      "-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  --------------------------------------------------------------------------------  \n",
      "Self CPU time total: 12.174ms\n",
      "Self CUDA time total: 474.717us\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(prof.key_averages(group_by_input_shape=True).table(sort_by=\"cpu_time_total\", row_limit=10))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch_profiler",
   "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.9.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}