4.2.3.1. Floating Point Model Preparation

A floating-point model consists of layers or modules that process on data. The torch.nn module provides all the network blocks for building floating-point models. All network modules in PyTorch are inherited from the torch.nn.Module. A network model itself is a network module composed of various smaller network modules. This embedded structure makes it easier for users to build and manage complex network architectures. Users can directly use the network modules provided by PyTorch to build floating-point models. Since the quantization is module-based, it is necessary to make some modifications to the model definition before the quantization, in the following aspects:

4.2.3.1.1. Operator Substitution

To convert a floating-point model to a quantized model, the functional operation that needs to quantize the output is converted into a module. (e.g., use torch.nn.ReLU to replace torch.nn.functional.relu). See the “Supported Operators” section in the API REFERENCE of the current document for the specific operators to be replaced.

4.2.3.1.2. Insert Quantization and Dequantization Nodes

For subsequent QAT and fixed-point prediction, it is necessary to insert a quantization node before the model’s input node and insert a dequantization nodes after the model’s output node. In terms of implementation, the quantization model as a whole starts with QuantStub and ends with DeQuantStub. But if the last layer outputs dequantized data such as class_idx (represented by the Tensor type in the QAT model instead of QTensor), DeQuantStub is not required. The following table lists the operators that do not require DeQuantStub when used as the output layer of the model.

The operator of DeQuantStub is not required as the model output layer
torch.Tensor.argmax / torch.argmax
horizon_plugin_pytorch.functional.argmax
horizon_plugin_pytorch.functional.filter
torch.max (the part where the return value is index does not need inverse quantization)

4.2.3.1.3. Set up Quantization Parameters

Specify which parts of the model need to be quantized by assigning a value to the model’s qconfig property. For example, use model.conv1.qconfig = None to set the model.conv layer to not be quantized, or set model.linear1 to use custom_qconfig instead of the global qconfig through model1.linear1.qconfig = custom_qconfig .

4.2.3.1.4. Example of Custom Floating Point Model

import torch
import torch.nn.quantized as nnq
import horizon_plugin_pytorch as horizon
from torch.quantization import QConfig, DeQuantStub
from torch import nn
from horizon_plugin_pytorch.quantization import (
    fuse_known_modules
    QuantStub 
)


class ExampleNet(nn.Module):
    def __init__(self):
        super(ExampleNet, self).__init__()
        self.quant = QuantStub()
        self.conv = nn.Conv2d(in_channels=1, out_channels=1, kernel_size=3)
        self.bn = nn.BatchNorm2d(num_features=1)
        # The add operation must use a FloatFunctional
        self.add = nn.quantized.FloatFunctional()
        self.act = nn.ReLU()
        self.out_conv = nn.Conv2d(in_channels=1, out_channels=1, kernel_size=3)
        self.dequant = DeQuantStub()

    def forward(self, x):
        # The quantized model as a whole generally start with QuantStub and end with DeQuantStub.
        x = self.quant(x)
        x = self.conv(x)
        x = self.bn(x)
        # If you want to fuse the add to conv, the first input to the add must come from the conv to which you want to fuse.
        # Pay attention to the calling method of the add here. When using FloatFunctional, you must use its specific method, instead of using forward directly.
        x = self.add.add(x, y)
        x = self.act(x)
        x = self.out_conv(x)
        x = self.dequant(x)
        return x

    def fuse_model(self):
        from horizon_plugin_pytorch import quantization

        quantization.fuse_modules(
            self,
            ["conv", "bn", "add", "act"],
            inplace=True,
        )

    def set_qconfig(self):
        # You don't need to call the set_qconfig method of the submodule, the submodule without qconfig will automatically use the qconfig of the parent module.
        self.qconfig = horizon.quantization.get_default_qat_qconfig()
        # If the last output layer of the network is conv, it can be set to out_qconfig separately to get more accurate output.
        self.out_conv.qconfig = (
            horizon.quantization.get_default_qat_out_qconfig()
        )