Files
nvidia--tensorrt/tools/pytorch-quantization/examples/finetune_quant_resnet50.ipynb
T
wehub-resource-sync c8a779b1bb
Docker Image CI / build-ubuntu2004 (push) Waiting to run
chore: import upstream snapshot with attribution
2026-07-13 13:36:55 +08:00

606 lines
39 KiB
Plaintext

{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Following the calibrate_quant_resnet50 example, now we fine tune the model"
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"WARNING: Logging before flag parsing goes to stderr.\n",
"W0608 21:25:39.018203 140228493526848 tensor_quant.py:96] Meaning of axis has changed since v2.0. Make sure to update.\n",
"W0608 21:25:39.019082 140228493526848 tensor_quant.py:96] Meaning of axis has changed since v2.0. Make sure to update.\n",
"W0608 21:25:39.019555 140228493526848 tensor_quant.py:96] Meaning of axis has changed since v2.0. Make sure to update.\n",
"W0608 21:25:39.020030 140228493526848 tensor_quant.py:96] Meaning of axis has changed since v2.0. Make sure to update.\n",
"W0608 21:25:39.020492 140228493526848 tensor_quant.py:96] Meaning of axis has changed since v2.0. Make sure to update.\n",
"W0608 21:25:39.020947 140228493526848 tensor_quant.py:96] Meaning of axis has changed since v2.0. Make sure to update.\n",
"W0608 21:25:39.021392 140228493526848 tensor_quant.py:96] Meaning of axis has changed since v2.0. Make sure to update.\n"
]
}
],
"source": [
"import datetime\n",
"import os\n",
"import sys\n",
"import time\n",
"\n",
"import torch\n",
"import torch.utils.data\n",
"from torch import nn\n",
"\n",
"from tqdm import tqdm\n",
"\n",
"import torchvision\n",
"from torchvision import transforms\n",
"\n",
"from pytorch_quantization import nn as quant_nn\n",
"from pytorch_quantization import calib\n",
"from pytorch_quantization.tensor_quant import QuantDescriptor\n",
"\n",
"from absl import logging\n",
"logging.set_verbosity(logging.FATAL) # Disable logging as they are too noisy in notebook"
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
"# For simplicity, import train and eval functions from the train script from torchvision instead of copything them here\n",
"sys.path.append(\"/raid/skyw/models/torchvision/references/classification/\")\n",
"from train import evaluate, train_one_epoch, load_data"
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"QuantResNet(\n",
" (conv1): QuantConv2d(\n",
" 3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False\n",
" (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=2.6387 calibrator=MaxCalibrator(track_amax=False) quant)\n",
" (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0000, 0.7817](64) calibrator=MaxCalibrator(track_amax=False) quant)\n",
" )\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): Bottleneck(\n",
" (conv1): QuantConv2d(\n",
" 64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
" (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=2.9730 calibrator=MaxCalibrator(track_amax=False) quant)\n",
" (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0000, 0.7266](64) calibrator=MaxCalibrator(track_amax=False) quant)\n",
" )\n",
" (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" (conv2): QuantConv2d(\n",
" 64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
" (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.0971 calibrator=MaxCalibrator(track_amax=False) quant)\n",
" (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0000, 0.4679](64) calibrator=MaxCalibrator(track_amax=False) quant)\n",
" )\n",
" (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" (conv3): QuantConv2d(\n",
" 64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
" (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.3318 calibrator=MaxCalibrator(track_amax=False) quant)\n",
" (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0000, 0.3936](256) calibrator=MaxCalibrator(track_amax=False) quant)\n",
" )\n",
" (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" (relu): ReLU(inplace=True)\n",
" (downsample): Sequential(\n",
" (0): QuantConv2d(\n",
" 64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
" (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=2.9730 calibrator=MaxCalibrator(track_amax=False) quant)\n",
" (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0000, 0.9879](256) calibrator=MaxCalibrator(track_amax=False) quant)\n",
" )\n",
" (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" )\n",
" )\n",
" (1): Bottleneck(\n",
" (conv1): QuantConv2d(\n",
" 256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
" (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.4872 calibrator=MaxCalibrator(track_amax=False) quant)\n",
" (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0000, 0.2618](64) calibrator=MaxCalibrator(track_amax=False) quant)\n",
" )\n",
" (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" (conv2): QuantConv2d(\n",
" 64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
" (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.0466 calibrator=MaxCalibrator(track_amax=False) quant)\n",
" (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0564, 0.5201](64) calibrator=MaxCalibrator(track_amax=False) quant)\n",
" )\n",
" (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" (conv3): QuantConv2d(\n",
" 64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
" (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.5106 calibrator=MaxCalibrator(track_amax=False) quant)\n",
" (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0000, 0.2946](256) calibrator=MaxCalibrator(track_amax=False) quant)\n",
" )\n",
" (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" (relu): ReLU(inplace=True)\n",
" )\n",
" (2): Bottleneck(\n",
" (conv1): QuantConv2d(\n",
" 256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
" (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.5250 calibrator=MaxCalibrator(track_amax=False) quant)\n",
" (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0533, 0.1921](64) calibrator=MaxCalibrator(track_amax=False) quant)\n",
" )\n",
" (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" (conv2): QuantConv2d(\n",
" 64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
" (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=0.9980 calibrator=MaxCalibrator(track_amax=False) quant)\n",
" (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0810, 0.2856](64) calibrator=MaxCalibrator(track_amax=False) quant)\n",
" )\n",
" (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" (conv3): QuantConv2d(\n",
" 64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
" (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.6532 calibrator=MaxCalibrator(track_amax=False) quant)\n",
" (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0197, 0.2752](256) calibrator=MaxCalibrator(track_amax=False) quant)\n",
" )\n",
" (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" (relu): ReLU(inplace=True)\n",
" )\n",
" )\n",
" (layer2): Sequential(\n",
" (0): Bottleneck(\n",
" (conv1): QuantConv2d(\n",
" 256, 128, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
" (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.5499 calibrator=MaxCalibrator(track_amax=False) quant)\n",
" (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0670, 0.3532](128) calibrator=MaxCalibrator(track_amax=False) quant)\n",
" )\n",
" (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" (conv2): QuantConv2d(\n",
" 128, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False\n",
" (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.1606 calibrator=MaxCalibrator(track_amax=False) quant)\n",
" (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0694, 0.2993](128) calibrator=MaxCalibrator(track_amax=False) quant)\n",
" )\n",
" (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" (conv3): QuantConv2d(\n",
" 128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
" (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.1425 calibrator=MaxCalibrator(track_amax=False) quant)\n",
" (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0000, 0.3917](512) calibrator=MaxCalibrator(track_amax=False) quant)\n",
" )\n",
" (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" (relu): ReLU(inplace=True)\n",
" (downsample): Sequential(\n",
" (0): QuantConv2d(\n",
" 256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False\n",
" (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.5499 calibrator=MaxCalibrator(track_amax=False) quant)\n",
" (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0000, 0.5662](512) calibrator=MaxCalibrator(track_amax=False) quant)\n",
" )\n",
" (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" )\n",
" )\n",
" (1): Bottleneck(\n",
" (conv1): QuantConv2d(\n",
" 512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
" (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.4626 calibrator=MaxCalibrator(track_amax=False) quant)\n",
" (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0370, 0.2522](128) calibrator=MaxCalibrator(track_amax=False) quant)\n",
" )\n",
" (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" (conv2): QuantConv2d(\n",
" 128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
" (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=0.8304 calibrator=MaxCalibrator(track_amax=False) quant)\n",
" (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0142, 0.2998](128) calibrator=MaxCalibrator(track_amax=False) quant)\n",
" )\n",
" (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" (conv3): QuantConv2d(\n",
" 128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
" (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.1722 calibrator=MaxCalibrator(track_amax=False) quant)\n",
" (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0000, 0.3038](512) calibrator=MaxCalibrator(track_amax=False) quant)\n",
" )\n",
" (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" (relu): ReLU(inplace=True)\n",
" )\n",
" (2): Bottleneck(\n",
" (conv1): QuantConv2d(\n",
" 512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
" (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.4864 calibrator=MaxCalibrator(track_amax=False) quant)\n",
" (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0653, 0.2383](128) calibrator=MaxCalibrator(track_amax=False) quant)\n",
" )\n",
" (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" (conv2): QuantConv2d(\n",
" 128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
" (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=0.9450 calibrator=MaxCalibrator(track_amax=False) quant)\n",
" (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0646, 0.2556](128) calibrator=MaxCalibrator(track_amax=False) quant)\n",
" )\n",
" (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" (conv3): QuantConv2d(\n",
" 128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
" (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=0.8535 calibrator=MaxCalibrator(track_amax=False) quant)\n",
" (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0162, 0.3522](512) calibrator=MaxCalibrator(track_amax=False) quant)\n",
" )\n",
" (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" (relu): ReLU(inplace=True)\n",
" )\n",
" (3): Bottleneck(\n",
" (conv1): QuantConv2d(\n",
" 512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
" (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.5229 calibrator=MaxCalibrator(track_amax=False) quant)\n",
" (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0648, 0.2814](128) calibrator=MaxCalibrator(track_amax=False) quant)\n",
" )\n",
" (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" (conv2): QuantConv2d(\n",
" 128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
" (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=0.9247 calibrator=MaxCalibrator(track_amax=False) quant)\n",
" (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0595, 0.2210](128) calibrator=MaxCalibrator(track_amax=False) quant)\n",
" )\n",
" (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" (conv3): QuantConv2d(\n",
" 128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
" (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=0.9747 calibrator=MaxCalibrator(track_amax=False) quant)\n",
" (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0127, 0.2956](512) calibrator=MaxCalibrator(track_amax=False) quant)\n",
" )\n",
" (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" (relu): ReLU(inplace=True)\n",
" )\n",
" )\n",
" (layer3): Sequential(\n",
" (0): Bottleneck(\n",
" (conv1): QuantConv2d(\n",
" 512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
" (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.5941 calibrator=MaxCalibrator(track_amax=False) quant)\n",
" (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0842, 0.3425](256) calibrator=MaxCalibrator(track_amax=False) quant)\n",
" )\n",
" (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" (conv2): QuantConv2d(\n",
" 256, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False\n",
" (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.3565 calibrator=MaxCalibrator(track_amax=False) quant)\n",
" (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0544, 0.2008](256) calibrator=MaxCalibrator(track_amax=False) quant)\n",
" )\n",
" (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" (conv3): QuantConv2d(\n",
" 256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
" (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.0293 calibrator=MaxCalibrator(track_amax=False) quant)\n",
" (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0000, 0.3212](1024) calibrator=MaxCalibrator(track_amax=False) quant)\n",
" )\n",
" (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" (relu): ReLU(inplace=True)\n",
" (downsample): Sequential(\n",
" (0): QuantConv2d(\n",
" 512, 1024, kernel_size=(1, 1), stride=(2, 2), bias=False\n",
" (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.5941 calibrator=MaxCalibrator(track_amax=False) quant)\n",
" (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0000, 0.3460](1024) calibrator=MaxCalibrator(track_amax=False) quant)\n",
" )\n",
" (1): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" )\n",
" )\n",
" (1): Bottleneck(\n",
" (conv1): QuantConv2d(\n",
" 1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
" (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.3305 calibrator=MaxCalibrator(track_amax=False) quant)\n",
" (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0407, 0.2942](256) calibrator=MaxCalibrator(track_amax=False) quant)\n",
" )\n",
" (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" (conv2): QuantConv2d(\n",
" 256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
" (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=0.9844 calibrator=MaxCalibrator(track_amax=False) quant)\n",
" (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0432, 0.2634](256) calibrator=MaxCalibrator(track_amax=False) quant)\n",
" )\n",
" (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" (conv3): QuantConv2d(\n",
" 256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
" (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=0.9333 calibrator=MaxCalibrator(track_amax=False) quant)\n",
" (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0000, 0.4969](1024) calibrator=MaxCalibrator(track_amax=False) quant)\n",
" )\n",
" (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" (relu): ReLU(inplace=True)\n",
" )\n",
" (2): Bottleneck(\n",
" (conv1): QuantConv2d(\n",
" 1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
" (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.3388 calibrator=MaxCalibrator(track_amax=False) quant)\n",
" (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0469, 0.2715](256) calibrator=MaxCalibrator(track_amax=False) quant)\n",
" )\n",
" (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" (conv2): QuantConv2d(\n",
" 256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
" (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=0.8617 calibrator=MaxCalibrator(track_amax=False) quant)\n",
" (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0397, 0.2100](256) calibrator=MaxCalibrator(track_amax=False) quant)\n",
" )\n",
" (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" (conv3): QuantConv2d(\n",
" 256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
" (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=0.7507 calibrator=MaxCalibrator(track_amax=False) quant)\n",
" (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0000, 0.3538](1024) calibrator=MaxCalibrator(track_amax=False) quant)\n",
" )\n",
" (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" (relu): ReLU(inplace=True)\n",
" )\n",
" (3): Bottleneck(\n",
" (conv1): QuantConv2d(\n",
" 1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
" (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.3554 calibrator=MaxCalibrator(track_amax=False) quant)\n",
" (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0553, 0.2390](256) calibrator=MaxCalibrator(track_amax=False) quant)\n",
" )\n",
" (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" (conv2): QuantConv2d(\n",
" 256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
" (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=0.9257 calibrator=MaxCalibrator(track_amax=False) quant)\n",
" (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0455, 0.2792](256) calibrator=MaxCalibrator(track_amax=False) quant)\n",
" )\n",
" (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" (conv3): QuantConv2d(\n",
" 256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
" (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=0.8117 calibrator=MaxCalibrator(track_amax=False) quant)\n",
" (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0000, 0.3126](1024) calibrator=MaxCalibrator(track_amax=False) quant)\n",
" )\n",
" (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" (relu): ReLU(inplace=True)\n",
" )\n",
" (4): Bottleneck(\n",
" (conv1): QuantConv2d(\n",
" 1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
" (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.4199 calibrator=MaxCalibrator(track_amax=False) quant)\n",
" (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0598, 0.2722](256) calibrator=MaxCalibrator(track_amax=False) quant)\n",
" )\n",
" (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" (conv2): QuantConv2d(\n",
" 256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
" (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=0.9274 calibrator=MaxCalibrator(track_amax=False) quant)\n",
" (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0459, 0.1919](256) calibrator=MaxCalibrator(track_amax=False) quant)\n",
" )\n",
" (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" (conv3): QuantConv2d(\n",
" 256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
" (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=0.8702 calibrator=MaxCalibrator(track_amax=False) quant)\n",
" (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0160, 0.3161](1024) calibrator=MaxCalibrator(track_amax=False) quant)\n",
" )\n",
" (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" (relu): ReLU(inplace=True)\n",
" )\n",
" (5): Bottleneck(\n",
" (conv1): QuantConv2d(\n",
" 1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
" (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.4258 calibrator=MaxCalibrator(track_amax=False) quant)\n",
" (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0618, 0.3995](256) calibrator=MaxCalibrator(track_amax=False) quant)\n",
" )\n",
" (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" (conv2): QuantConv2d(\n",
" 256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
" (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.2256 calibrator=MaxCalibrator(track_amax=False) quant)\n",
" (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0498, 0.2236](256) calibrator=MaxCalibrator(track_amax=False) quant)\n",
" )\n",
" (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" (conv3): QuantConv2d(\n",
" 256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
" (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.3560 calibrator=MaxCalibrator(track_amax=False) quant)\n",
" (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0180, 0.3288](1024) calibrator=MaxCalibrator(track_amax=False) quant)\n",
" )\n",
" (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" (relu): ReLU(inplace=True)\n",
" )\n",
" )\n",
" (layer4): Sequential(\n",
" (0): Bottleneck(\n",
" (conv1): QuantConv2d(\n",
" 1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
" (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.3915 calibrator=MaxCalibrator(track_amax=False) quant)\n",
" (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0744, 0.3415](512) calibrator=MaxCalibrator(track_amax=False) quant)\n",
" )\n",
" (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" (conv2): QuantConv2d(\n",
" 512, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False\n",
" (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.1571 calibrator=MaxCalibrator(track_amax=False) quant)\n",
" (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0432, 0.3993](512) calibrator=MaxCalibrator(track_amax=False) quant)\n",
" )\n",
" (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" (conv3): QuantConv2d(\n",
" 512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
" (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.1295 calibrator=MaxCalibrator(track_amax=False) quant)\n",
" (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0198, 0.3546](2048) calibrator=MaxCalibrator(track_amax=False) quant)\n",
" )\n",
" (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" (relu): ReLU(inplace=True)\n",
" (downsample): Sequential(\n",
" (0): QuantConv2d(\n",
" 1024, 2048, kernel_size=(1, 1), stride=(2, 2), bias=False\n",
" (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.3915 calibrator=MaxCalibrator(track_amax=False) quant)\n",
" (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0197, 0.6413](2048) calibrator=MaxCalibrator(track_amax=False) quant)\n",
" )\n",
" (1): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" )\n",
" )\n",
" (1): Bottleneck(\n",
" (conv1): QuantConv2d(\n",
" 2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
" (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=3.9348 calibrator=MaxCalibrator(track_amax=False) quant)\n",
" (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0482, 0.7003](512) calibrator=MaxCalibrator(track_amax=False) quant)\n",
" )\n",
" (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" (conv2): QuantConv2d(\n",
" 512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
" (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.1277 calibrator=MaxCalibrator(track_amax=False) quant)\n",
" (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0379, 0.2257](512) calibrator=MaxCalibrator(track_amax=False) quant)\n",
" )\n",
" (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" (conv3): QuantConv2d(\n",
" 512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
" (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=0.8992 calibrator=MaxCalibrator(track_amax=False) quant)\n",
" (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0211, 0.2427](2048) calibrator=MaxCalibrator(track_amax=False) quant)\n",
" )\n",
" (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" (relu): ReLU(inplace=True)\n",
" )\n",
" (2): Bottleneck(\n",
" (conv1): QuantConv2d(\n",
" 2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
" (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=5.2181 calibrator=MaxCalibrator(track_amax=False) quant)\n",
" (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0601, 0.4541](512) calibrator=MaxCalibrator(track_amax=False) quant)\n",
" )\n",
" (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" (conv2): QuantConv2d(\n",
" 512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
" (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.2051 calibrator=MaxCalibrator(track_amax=False) quant)\n",
" (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0338, 0.1416](512) calibrator=MaxCalibrator(track_amax=False) quant)\n",
" )\n",
" (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" (conv3): QuantConv2d(\n",
" 512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
" (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.1045 calibrator=MaxCalibrator(track_amax=False) quant)\n",
" (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0119, 0.2798](2048) calibrator=MaxCalibrator(track_amax=False) quant)\n",
" )\n",
" (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" (relu): ReLU(inplace=True)\n",
" )\n",
" )\n",
" (avgpool): AdaptiveAvgPool2d(output_size=(1, 1))\n",
" (fc): QuantLinear(\n",
" in_features=2048, out_features=1000, bias=True\n",
" (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=5.5345 calibrator=MaxCalibrator(track_amax=False) quant)\n",
" (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.1716, 0.7371](1000) calibrator=MaxCalibrator(track_amax=False) quant)\n",
" )\n",
")"
]
},
"execution_count": 3,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"from pytorch_quantization import quant_modules\n",
"quant_modules.initialize()\n",
"\n",
"# Create and load the calibrated model\n",
"model = torchvision.models.resnet50()\n",
"model.load_state_dict(torch.load(\"/tmp/quant_resnet50-calibrated.pth\"))\n",
"model.cuda()"
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [],
"source": [
"data_path = \"/raid/data/imagenet/imagenet_pytorch\"\n",
"\n",
"traindir = os.path.join(data_path, 'train')\n",
"valdir = os.path.join(data_path, 'val')\n",
"dataset, dataset_test, train_sampler, test_sampler = load_data(traindir, valdir, False, False)\n",
"\n",
"data_loader = torch.utils.data.DataLoader(\n",
" dataset, batch_size=256,\n",
" sampler=train_sampler, num_workers=4, pin_memory=True)\n",
"\n",
"data_loader_test = torch.utils.data.DataLoader(\n",
" dataset_test, batch_size=256,\n",
" sampler=test_sampler, num_workers=4, pin_memory=True)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Quantized fine tuning\n",
"Let's fine tune the model with fake quantization. We only fine tune for 1 epoch as an example."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"criterion = nn.CrossEntropyLoss()\n",
"optimizer = torch.optim.SGD(model.parameters(), lr=0.0001)\n",
"lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.1)\n",
"\n",
"data_loader = torch.utils.data.DataLoader(\n",
" dataset, batch_size=128,\n",
" sampler=train_sampler, num_workers=16, pin_memory=True)\n",
"\n",
"# Training takes about one and half hour per epoch on single V100\n",
"train_one_epoch(model, criterion, optimizer, data_loader, \"cuda\", 0, 100)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Evaluate the fine tuned model"
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {
"scrolled": true
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Test: [ 0/196] eta: 0:09:36 loss: 0.4680 (0.4680) acc1: 85.9375 (85.9375) acc5: 98.0469 (98.0469) time: 2.9406 data: 2.1852 max mem: 16096\n",
"Test: [ 10/196] eta: 0:01:58 loss: 0.6694 (0.6522) acc1: 83.2031 (82.9545) acc5: 96.0938 (96.1293) time: 0.6346 data: 0.1988 max mem: 16096\n",
"Test: [ 20/196] eta: 0:01:30 loss: 0.6738 (0.6733) acc1: 82.0312 (82.4777) acc5: 95.7031 (95.7961) time: 0.3928 data: 0.0001 max mem: 16096\n",
"Test: [ 30/196] eta: 0:01:18 loss: 0.6219 (0.6322) acc1: 84.3750 (83.9718) acc5: 95.7031 (96.0181) time: 0.3859 data: 0.0001 max mem: 16096\n",
"Test: [ 40/196] eta: 0:01:10 loss: 0.6801 (0.6750) acc1: 81.6406 (82.5934) acc5: 95.7031 (95.9604) time: 0.3861 data: 0.0001 max mem: 16096\n",
"Test: [ 50/196] eta: 0:01:04 loss: 0.6937 (0.6724) acc1: 80.0781 (82.3529) acc5: 96.8750 (96.0938) time: 0.3834 data: 0.0001 max mem: 16096\n",
"Test: [ 60/196] eta: 0:00:58 loss: 0.7149 (0.6849) acc1: 80.0781 (81.9864) acc5: 96.4844 (96.1066) time: 0.3854 data: 0.0001 max mem: 16096\n",
"Test: [ 70/196] eta: 0:00:53 loss: 0.6616 (0.6716) acc1: 80.8594 (82.2843) acc5: 96.4844 (96.1983) time: 0.3859 data: 0.0001 max mem: 16096\n",
"Test: [ 80/196] eta: 0:00:48 loss: 0.6510 (0.6968) acc1: 81.2500 (81.7467) acc5: 95.7031 (95.9201) time: 0.3860 data: 0.0001 max mem: 16096\n",
"Test: [ 90/196] eta: 0:00:44 loss: 0.9469 (0.7444) acc1: 76.1719 (80.6834) acc5: 92.5781 (95.4370) time: 0.3868 data: 0.0001 max mem: 16096\n",
"Test: [100/196] eta: 0:00:39 loss: 1.1594 (0.7964) acc1: 70.7031 (79.5521) acc5: 90.6250 (94.8755) time: 0.3864 data: 0.0001 max mem: 16096\n",
"Test: [110/196] eta: 0:00:35 loss: 1.1594 (0.8214) acc1: 72.2656 (79.0365) acc5: 91.4062 (94.6298) time: 0.3836 data: 0.0001 max mem: 16096\n",
"Test: [120/196] eta: 0:00:31 loss: 0.9820 (0.8389) acc1: 76.1719 (78.7771) acc5: 92.1875 (94.3634) time: 0.3856 data: 0.0001 max mem: 16096\n",
"Test: [130/196] eta: 0:00:26 loss: 1.0825 (0.8705) acc1: 72.6562 (77.9610) acc5: 91.4062 (94.0303) time: 0.3866 data: 0.0001 max mem: 16096\n",
"Test: [140/196] eta: 0:00:22 loss: 1.1088 (0.8889) acc1: 72.2656 (77.6125) acc5: 91.4062 (93.8137) time: 0.3879 data: 0.0001 max mem: 16096\n",
"Test: [150/196] eta: 0:00:18 loss: 1.1069 (0.9059) acc1: 73.0469 (77.2998) acc5: 91.0156 (93.5586) time: 0.3914 data: 0.0002 max mem: 16096\n",
"Test: [160/196] eta: 0:00:14 loss: 1.1360 (0.9197) acc1: 73.0469 (77.0380) acc5: 90.2344 (93.3472) time: 0.3898 data: 0.0002 max mem: 16096\n",
"Test: [170/196] eta: 0:00:10 loss: 1.2171 (0.9371) acc1: 71.8750 (76.6265) acc5: 89.8438 (93.1721) time: 0.3845 data: 0.0002 max mem: 16096\n",
"Test: [180/196] eta: 0:00:06 loss: 1.2493 (0.9527) acc1: 68.7500 (76.2992) acc5: 90.2344 (93.0205) time: 0.3815 data: 0.0001 max mem: 16096\n",
"Test: [190/196] eta: 0:00:02 loss: 1.0816 (0.9511) acc1: 71.4844 (76.3089) acc5: 92.1875 (93.0465) time: 0.3736 data: 0.0001 max mem: 16096\n",
"Test: Total time: 0:01:17\n",
" * Acc@1 76.426 Acc@5 93.080\n"
]
}
],
"source": [
"with torch.no_grad():\n",
" evaluate(model, criterion, data_loader_test, device=\"cuda\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"After only 1 epoch of quantized fine tuning, top-1 improved from ~76.1 to 76.426. Train longer with lr anealing can improve accuracy further"
]
}
],
"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.10"
}
},
"nbformat": 4,
"nbformat_minor": 2
}