7.3.2. Operator Fusion

7.3.2.1. Why Do Operator Fusion

Operator fusion can both speed up the computation and improve the quantization accuracy.

7.3.2.2. Speed up Computation

For example, fuse Conv and ReLU into ConvReLU2d , as shown in the figure below: ReLU reads the calculation result of Conv and then calculates it, while ConvReLU2d directly calculates the result of Conv , which saves the process of reading data and speeds up the calculation.

This is a simple example, and it will be much more complicated than this in actual situation.

7.3.2.3. Improve Quantization Accuracy

If each operator is computed independently, then the output is 8 bit data. If operator fusion is used, then the Conv0 output is 32 bit data, as shown in the figure below:

7.3.2.4. Fusion Operator

from torch import nn
import torch.nn.quantized as nnq

# The following operator fusions are currently supported:
(nn.Conv2d, nn.BatchNorm2d, nn.ReLU)
(nn.Conv2d, nn.ReLU)
(nn.Conv2d, nn.BatchNorm2d, nnq.FloatFunctional)
(nn.Conv2d, nn.BatchNorm2d, nnq.FloatFunctional, nn.ReLU)
(nn.Conv2d, nnq.FloatFunctional)
(nn.Conv2d, nnq.FloatFunctional, nn.ReLU)
(nn.ConvTranspose2d, nn.ReLU)
(nn.ConvTranspose2d, nnq.FloatFunctional)
(nn.ConvTranspose2d, nnq.FloatFunctional, nn.ReLU)
(nn.ConvTranspose2d, nn.BatchNorm2d)
(nn.ConvTranspose2d, nn.BatchNorm2d, nn.ReLU)
(nn.ConvTranspose2d, nn.BatchNorm2d, nnq.FloatFunctional)
(nn.ConvTranspose2d, nn.BatchNorm2d, nnq.FloatFunctional, nn.ReLU)
(nn.Conv2d, nn.BatchNorm2d, nn.ReLU6)
(nn.Conv2d, nn.ReLU6)
(nn.Conv2d, nn.BatchNorm2d, nnq.FloatFunctional, nn.ReLU6)
(nn.Conv2d, nnq.FloatFunctional, nn.ReLU6)
(nn.ConvTranspose2d, nn.ReLU6)
(nn.ConvTranspose2d, nnq.FloatFunctional, nn.ReLU6)
(nn.ConvTranspose2d, nn.BatchNorm2d, nn.ReLU6)
(nn.ConvTranspose2d, nn.BatchNorm2d, nnq.FloatFunctional, nn.ReLU6)

7.3.2.5. Purpose of BN (Batch Normalization) Absorption

BN absorption aims to reduce the computation of the deployed model. Since BN is a linear transformation process, the transformation parameters of BN can be absorbed into the Conv parameters when BN and Conv appear together, thereby eliminating the computation of BN in the deployed model.

Convert Conv2d + BN2d to Conv2d by absorbing BN .

7.3.2.6. BN Absorption Method

Currently the tool supports the Conv -> BN mode to absorb BN .

The absorption method is as follows:

7.3.2.7. Operator Fusion Example

Example 1: Take operator subscripts for fusion.

import torch
import horizon_plugin_pytorch as D-Robotics
from torch.quantization import DeQuantStub
from horizon_plugin_pytorch.quantization import QuantStub

class ModelForFusion(torch.nn.Sequential):
    def __init__(
        self,
        in_channels,
        out_channels,
        kernel_size,
        dequant_out=False,
    ):
        super(ModelForFusion, self).__init__(
            QuantStub(),
            nn.Conv2d(
                in_channels,
                out_channels,
                kernel_size,
            ),
            nn.BatchNorm2d(num_features=out_channels),
            DeQuantStub() if dequant_out else nn.Identity(),
        )

float_net = ModelForFusion(
    1,
    2,
    1,
)
# Since the network to be fused is inherited from torch.nn.Sequential, take the subscripts of conv and bn in the network and put them in the list to determine the operators that need to be fused.
horizon.quantization.fuse_modules(
    float_net, ["1", "2"], inplace=True
)

Example 2: Fusion with operator names in a list.

from torch import nn


class ModelForFusion(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(2, 2, 1, bias=None)
        self.bn = nn.BatchNorm2d(2)
        self.relu = nn.ReLU(inplace=True)
        self.quant = QuantStub()
        self.dequant = DeQuantStub()

    def forward(self, x):
        x_i = self.quant(x)
        x = self.conv(x_i)
        x = self.bn(x)
        x = self.relu(x)
        x = self.dequant(x)
        return x

model = ModelForFusion().train()

Since the network is inherited from Module, each operator in the network has a variable name, and the variable names of conv and bn in the network are taken into the list for fusion
horizon.quantization.fuse_modules(
    model,
    ["conv", "bn", "relu"],
    inplace=True,
)