PyTorch Version (vai_q_pytorch)

vai_q_pytorch is designed to work as a PyTorch plugin. Xilinx provides the simplest APIs to introduce the FPGA-friendly quantization feature. For a well-defined model, you only need to add a few lines to get a quantize model object. To do so, follow these steps:

You can use module partial quantization if not all the sub-modules in a model need to be quantized. Besides using general vai_q_pytorch APIs, the QuantStub/DeQuantStub operator pair can be used to realize it. The following example demonstrates how to quantize subm0 and subm2, but not quantize subm1.

from pytorch_nndct.nn import QuantStub, DeQuantStub

class WholeModule(torch.nn.module):
    def __init__(self,...):
        self.subm0 = ...
        self.subm1 = ...
        self.subm2 = ...

        # define QuantStub/DeQuantStub submodules
        self.quant = QuantStub()
        self.dequant = DeQuantStub()

    def forward(self, input):
        input = self.quant(input) # begin of part to be quantized
        output0 = self.subm0(input)
        output0 = self.dequant(output0) # end of part to be quantized

        output1 = self.subm1(output0)

        output1 = self.quant(output1) # begin of part to be quantized
        output2 = self.subm2(output1)
        output2 = self.dequant(output2) # end of part to be quantized

Generally, there is a small accuracy loss after quantization, but for some networks such as MobileNets, the accuracy loss can be large. In this situation, first try fast finetune. If fast finetune still does not yield satisfactory results, quantize finetuning can be used to further improve the accuracy of the quantized models.

The AdaQuant algorithm¹ uses a small set of unlabeled data. It not only calibrates the activations but also finetunes the weights. The Vitis AI quantizer implements this algorithm and call it "fast finetuning" or "advanced calibration." Though slightly slower, fast finetuning can achieve better performance than quantize calibration. Similar to quantize finetuning, each run of fast finetuning produces a different result.

Fast finetuning does not train the model, and only needs a limited number of iterations. For classification models on Imagenet dataset, 1000 images are enough. Fast finetuning only needs some modification based on the model evaluation script. There is no need to set up the optimizer for training. To use fast finetuning, a function for model forwarding iteration is needed and will be called during fast finetuning. Re-calibration with the original inference code is recommended.

You can find a complete example in the open source example.

# fast finetune model or load finetuned parameter before test 
  if fast_finetune == True:
      ft_loader, _ = load_data(
          subset_len=1024,
          train=False,
          batch_size=batch_size,
          sample_method=None,
          data_dir=args.data_dir,
          model_name=model_name)
      if quant_mode == 'calib':
          quantizer.fast_finetune(evaluate, (quant_model, ft_loader, loss_fn))
      elif quant_mode == 'test':
          quantizer.load_ft_param()

python resnet18_quant.py --quant_mode calib --fast_finetune

python resnet18_quant.py --quant_mode test --fast_finetune

Assuming that there is a pre-defined model architecture, use the following steps to do quantization aware training. Take the ResNet18 model from Torchvision as an example. The complete model definition is here.

1. Check if there are non-module operations to be quantized. ResNet18 uses ‘+’ to add two tensors. Replace them with pytorch_nndct.nn.modules.functional.Add.
2. Check if there are modules to be called multiple times. Usually such modules have no weights; the most common one is the torch.nn.ReLu module. Define multiple such modules and then call them separately in a forward pass. The revised definition that meets the requirements is as follows:

   class BasicBlock(nn.Module):
     expansion = 1
   
     def __init__(self,
                  inplanes,
                  planes,
                  stride=1,
                  downsample=None,
                  groups=1,
                  base_width=64,
                  dilation=1,
                  norm_layer=None):
       super(BasicBlock, self).__init__()
       if norm_layer is None:
         norm_layer = nn.BatchNorm2d
       if groups != 1 or base_width != 64:
         raise ValueError('BasicBlock only supports groups=1 and base_width=64')
       if dilation > 1:
         raise NotImplementedError("Dilation > 1 not supported in BasicBlock")
       # Both self.conv1 and self.downsample layers downsample the input when stride != 1
       self.conv1 = conv3x3(inplanes, planes, stride)
       self.bn1 = norm_layer(planes)
       self.relu1 = nn.ReLU(inplace=True)
       self.conv2 = conv3x3(planes, planes)
       self.bn2 = norm_layer(planes)
       self.downsample = downsample
       self.stride = stride
   
       # Use a functional module to replace ‘+’
       self.skip_add = functional.Add()
   
       # Additional defined module
       self.relu2 = nn.ReLU(inplace=True)
   
     def forward(self, x):
       identity = x
   
       out = self.conv1(x)
       out = self.bn1(out)
       out = self.relu1(out)
   
       out = self.conv2(out)
       out = self.bn2(out)
   
       if self.downsample is not None:
         identity = self.downsample(x)
       
       # Use function module instead of ‘+’
       # out += identity
       out = self.skip_add(out, identity)
       out = self.relu2(out)
   
       return out
   

3. Insert QuantStub and DeQuantStub.

   Use QuantStub to quantize the inputs of the network and DeQuantStub to de-quantize the outputs of the network. Any sub-network from QuantStub to DeQuantStub in a forward pass will be quantized. Multiple QuantStub-DeQuantStub pairs are allowed.

   class ResNet(nn.Module):
   
     def __init__(self,
                  block,
                  layers,
                  num_classes=1000,
                  zero_init_residual=False,
                  groups=1,
                  width_per_group=64,
                  replace_stride_with_dilation=None,
                  norm_layer=None):
       super(ResNet, self).__init__()
       if norm_layer is None:
         norm_layer = nn.BatchNorm2d
       self._norm_layer = norm_layer
   
       self.inplanes = 64
       self.dilation = 1
       if replace_stride_with_dilation is None:
         # each element in the tuple indicates if we should replace
         # the 2x2 stride with a dilated convolution instead
         replace_stride_with_dilation = [False, False, False]
       if len(replace_stride_with_dilation) != 3:
         raise ValueError(
             "replace_stride_with_dilation should be None "
             "or a 3-element tuple, got {}".format(replace_stride_with_dilation))
       self.groups = groups
       self.base_width = width_per_group
       self.conv1 = nn.Conv2d(
           3, self.inplanes, kernel_size=7, stride=2, padding=3, bias=False)
       self.bn1 = norm_layer(self.inplanes)
       self.relu = nn.ReLU(inplace=True)
       self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
       self.layer1 = self._make_layer(block, 64, layers[0])
       self.layer2 = self._make_layer(
           block, 128, layers[1], stride=2, dilate=replace_stride_with_dilation[0])
       self.layer3 = self._make_layer(
           block, 256, layers[2], stride=2, dilate=replace_stride_with_dilation[1])
       self.layer4 = self._make_layer(
           block, 512, layers[3], stride=2, dilate=replace_stride_with_dilation[2])
       self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
       self.fc = nn.Linear(512 * block.expansion, num_classes)
   
       self.quant_stub = nndct_nn.QuantStub()
       self.dequant_stub = nndct_nn.DeQuantStub()
   
       for m in self.modules():
         if isinstance(m, nn.Conv2d):
           nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
         elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
           nn.init.constant_(m.weight, 1)
           nn.init.constant_(m.bias, 0)
   
       # Zero-initialize the last BN in each residual branch,
       # so that the residual branch starts with zeros, and each residual block behaves like an identity.
       # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
       if zero_init_residual:
         for m in self.modules():
           if isinstance(m, Bottleneck):
             nn.init.constant_(m.bn3.weight, 0)
           elif isinstance(m, BasicBlock):
             nn.init.constant_(m.bn2.weight, 0)
   
     def forward(self, x):
       x = self.quant_stub(x)
   
       x = self.conv1(x)
       x = self.bn1(x)
       x = self.relu(x)
       x = self.maxpool(x)
   
       x = self.layer1(x)
       x = self.layer2(x)
       x = self.layer3(x)
       x = self.layer4(x)
   
       x = self.avgpool(x)
       x = torch.flatten(x, 1)
       x = self.fc(x)
       x = self.dequant_stub(x)
       return x
   

4. Use QAT APIs to create the quantizer and train the model.

   def _resnet(arch, block, layers, pretrained, progress, **kwargs):
     model = ResNet(block, layers, **kwargs)
     if pretrained:
       #state_dict = load_state_dict_from_url(model_urls[arch], progress=progress)
       state_dict = torch.load(model_urls[arch])
       model.load_state_dict(state_dict)
     return model
   
   def resnet18(pretrained=False, progress=True, **kwargs):
     r"""ResNet-18 model from
       `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>'_
   
       Args:
           pretrained (bool): If True, returns a model pre-trained on ImageNet
           progress (bool): If True, displays a progress bar of the download to stderr
       """
     return _resnet('resnet18', BasicBlock, [2, 2, 2, 2], pretrained, progress,
                    **kwargs)
   
   model = resnet18(pretrained=True)
   
   # Generate dummy inputs.
   input = torch.randn([batch_size, 3, 224, 224], dtype=torch.float32)
   
   # Create a quantizer
   quantizer = torch_quantizer(quant_mode = 'calib',
                              module = model, 
                              input_args = input,
                              bitwidth = 8,
                              qat_proc = True)
   quantized_model = quantizer.quant_model
   optimizer = torch.optim.Adam(
           quantized_model.parameters(), 
           lr, 
           weight_decay=weight_decay)
   
   # Use the optimizer to train the model, just like a normal float model.
   …
   

5. Convert the trained model to a deployable model.

   After training, dump the quantized model to xmodel. (batch size=1 is must for compilation of xmodel).

   # vai_q_pytorch interface function: deploy the trained model and convert xmodel
     # need at least 1 iteration of inference with batch_size=1 
     quantizer.deploy(quantized_model)
     deployable_model = quantizer.deploy_model
     val_dataset2 = torch.utils.data.Subset(val_dataset, list(range(1)))
     val_loader2 = torch.utils.data.DataLoader(
         val_dataset,
         batch_size=1,
         shuffle=False,
         num_workers=workers,
         pin_memory=True)
     validate(val_loader2, deployable_model, criterion, gpu)
     quantizer.export_xmodel()
   

This section introduces the usage of execution tools and APIs to implement quantization and generate a model to be deployed on the target hardware. The APIs in the module pytorch_binding/pytorch_nndct/apis/quant_api.py are as follows: