算子融合

训练工具支持的算子融合可分为两大类：1. 吸收 BN；2. 融合 Add、ReLU(6)。

吸收 BN

吸收 BN 的目的是为了减少模型的计算量。因为 BN 是线性变换过程，因此，当 BN 和 Conv 一起出现的时候，可以把 BN 的参数吸收到 Conv 的参数中，从而在部署的模型中消除 BN 的计算。

吸收的计算过程如下：

通过吸收 BN ，可以把 Conv2d + BN2d 简化为 Conv2d：

融合 Add、ReLU(6)

和 CUDA Kernel Fusion 中将 CUDA Kernel 融合以提高计算速度不同，训练工具支持的融合更加偏重量化层面。

BPU 硬件针对常见的模型基本结构做了优化，在计算 Conv -> Add -> ReLU 这种算子组合时，可使算子间的数据传递保留高精度的状态，提高模型整体的数值精度。因此在对模型进行量化时，我们可以将 Conv -> Add -> ReLU 视为一个整体。

算子融合除了可以使中间结果保留高精度状态之外，也可以省去将中间结果转化为低精度表示的过程，因此执行速度和不融合相比也会更快。

由于算子融合既可以提高模型精度，又可以提高模型速度，一般应该对所有可融合的部分进行融合。图模式下，融合过程是工具自动进行的，您可不过多关注。

实现原理

得益于 FX 可以获取计算图的优势，训练工具可以自动化地对模型的计算图进行分析，根据预定义的 fusion pattern 对可融合部分进行匹配，并通过 submodule 替换实现融合的操作。

融合 Conv+BN

import torch
from torch import nn
from torch.quantization import DeQuantStub

from horizon_plugin_pytorch import qint8, set_march
from horizon_plugin_pytorch.quantization import QuantStub, get_qconfig, prepare
from horizon_plugin_pytorch.quantization.qconfig_setter import *


class ModelForFusion(torch.nn.Module):
    def __init__(self):
        super(ModelForFusion, self).__init__()
        self.quant = QuantStub()
        self.conv = nn.Conv2d(3, 3, 3)
        self.bn = nn.BatchNorm2d(3)
        self.dequant = DeQuantStub()

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

        return x


float_model = ModelForFusion()

set_march("nash-m")
fused_model = prepare(
    float_model,
    torch.rand(1, 3, 32, 32),
    QconfigSetter(
        get_qconfig(), [ModuleNameTemplate({"": qint8}), ConvDtypeTemplate()]
    ),
)

print(fused_model)
"""
GraphModuleImpl(
  (quant): QuantStub(
    (activation_post_process): FakeQuantize(
      dtype=qint8, fake_quant_enabled=True, observer_enabled=True, qscheme=torch.per_tensor_symmetric, ch_axis=-1, scale=tensor([1.]), zero_point=tensor([0])
      (activation_post_process): MinMaxObserver(averaging_constant=0.01)
    )
  )
  (conv): Conv2d(
    3, 3, kernel_size=(3, 3), stride=(1, 1)
    (weight_fake_quant): FakeQuantize(
      dtype=qint8, fake_quant_enabled=True, observer_enabled=True, qscheme=torch.per_channel_symmetric, ch_axis=0, scale=tensor([1., 1., 1.]), zero_point=tensor([0, 0, 0])
      (activation_post_process): MinMaxObserver(averaging_constant=0.01)
    )
  )
  (bn): Identity()
  (dequant): DeQuantStub()
)
"""

可以看到，对模型执行算子融合操作后，BN 被吸收进 Conv 中，原本的 submodule 被替换为 Identity。

融合 Conv+Add/ReLU

import torch
from torch import nn
from torch.quantization import DeQuantStub

from horizon_plugin_pytorch import qint8, set_march
from horizon_plugin_pytorch.quantization import QuantStub, get_qconfig, prepare
from horizon_plugin_pytorch.quantization.qconfig_setter import *


class ModelForFusion(torch.nn.Module):
    def __init__(self):
        super(ModelForFusion, self).__init__()
        self.quantx = QuantStub()
        self.quanty = QuantStub()
        self.conv = nn.Conv2d(3, 3, 3)
        self.bn = nn.BatchNorm2d(3)
        self.relu = nn.ReLU()
        self.dequant = DeQuantStub()

    def forward(self, x, y):
        x = self.quantx(x)
        y = self.quanty(y)
        x = self.conv(x)
        x = self.bn(x)
        x = x + y
        x = self.relu(x)
        x = self.dequant(x)
        return x


float_model = ModelForFusion()

set_march("nash-m")
fused_model = prepare(
    float_model,
    (torch.rand(1, 3, 32, 32), torch.rand(1, 3, 30, 30)),
    QconfigSetter(
        get_qconfig(), [ModuleNameTemplate({"": qint8}), ConvDtypeTemplate()]
    ),
)

print(fused_model)
"""
GraphModuleImpl(
  (quantx): QuantStub(
    (activation_post_process): FakeQuantize(
      dtype=qint8, fake_quant_enabled=True, observer_enabled=True, qscheme=torch.per_tensor_symmetric, ch_axis=-1, scale=tensor([1.]), zero_point=tensor([0])
      (activation_post_process): MinMaxObserver(averaging_constant=0.01)
    )
  )
  (quanty): QuantStub(
    (activation_post_process): FakeQuantize(
      dtype=qint8, fake_quant_enabled=True, observer_enabled=True, qscheme=torch.per_tensor_symmetric, ch_axis=-1, scale=tensor([1.]), zero_point=tensor([0])
      (activation_post_process): MinMaxObserver(averaging_constant=0.01)
    )
  )
  (conv): Conv2d(
    3, 3, kernel_size=(3, 3), stride=(1, 1)
    (weight_fake_quant): FakeQuantize(
      dtype=qint8, fake_quant_enabled=True, observer_enabled=True, qscheme=torch.per_channel_symmetric, ch_axis=0, scale=tensor([1., 1., 1.]), zero_point=tensor([0, 0, 0])
      (activation_post_process): MinMaxObserver(averaging_constant=0.01)
    )
  )
  (bn): Identity()
  (relu): ReLU()
  (dequant): DeQuantStub()
  (_generated_add_0): FloatFunctional()
)
"""

工具自动地将模型中 x = x + y 的加号替换为了名为 _generated_add_0 的 Module 形式，以支持算子融合和量化的相关操作。

Each layer out qconfig:
+------------------+----------------------------------------------------------------------+--------------------------+-----------------+-------------------+-----------------------------------------+
| Module Name      | Module Type                                                          | Input dtype              | out dtype       | ch_axis           | observer                                |
|------------------+----------------------------------------------------------------------+--------------------------+-----------------+-------------------+-----------------------------------------|
| quantx           | horizon_plugin_pytorch.nn.qat.stubs.QuantStub                        | [torch.float32]          | ['qint8']       | -1                | MinMaxObserver(averaging_constant=0.01) |
| quanty           | horizon_plugin_pytorch.nn.qat.stubs.QuantStub                        | [torch.float32]          | ['qint8']       | -1                | MinMaxObserver(averaging_constant=0.01) |
| conv             | horizon_plugin_pytorch.nn.qat.conv2d.Conv2d                          | ['qint8']                | [torch.float32] | activation = None |                                         |
| _generated_add_0 | horizon_plugin_pytorch.nn.qat.functional_modules.FloatFunctional.add | [torch.float32, 'qint8'] | [torch.float32] | activation = None |                                         |
| relu             | torch.nn.modules.activation.ReLU                                     | [torch.float32]          | [torch.float32] | qconfig = None    |                                         |
| dequant          | horizon_plugin_pytorch.nn.qat.stubs.DeQuantStub                      | [torch.float32]          | [torch.float32] | qconfig = None    |                                         |
+------------------+----------------------------------------------------------------------+--------------------------+-----------------+-------------------+-----------------------------------------+

从模型的检查结果中可以看到，量化阶段 conv 和 add, relu 之间使用了 float32 类型以模拟融合后的“中间高精度”，在后续的模型部署阶段会将多个计算进行合并，实现融合。

版本迁移说明

上述文档展示了使用 QconfigSetter 进行量化配置时的融合过程，在此之前工具提供了另一套 TemplateQconfigSetter 的配置方式

在旧的配置方式下，融合结果为：

GraphModuleImpl(
  (quantx): QuantStub(
    (activation_post_process): FakeQuantize(
      dtype=qint8, fake_quant_enabled=True, observer_enabled=True, qscheme=torch.per_tensor_symmetric, ch_axis=-1, scale=tensor([1.]), zero_point=tensor([0])
      (activation_post_process): MinMaxObserver(averaging_constant=0.01)
    )
  )
  (quanty): QuantStub(
    (activation_post_process): FakeQuantize(
      dtype=qint8, fake_quant_enabled=True, observer_enabled=True, qscheme=torch.per_tensor_symmetric, ch_axis=-1, scale=tensor([1.]), zero_point=tensor([0])
      (activation_post_process): MinMaxObserver(averaging_constant=0.01)
    )
  )
  (conv): Identity()
  (bn): Identity()
  (relu): Identity()
  (dequant): DeQuantStub()
  (_generated_add_0): ConvAddReLU2d(
    3, 3, kernel_size=(3, 3), stride=(1, 1)
    (weight_fake_quant): FakeQuantize(
      dtype=qint8, fake_quant_enabled=True, observer_enabled=True, qscheme=torch.per_channel_symmetric, ch_axis=0, scale=tensor([1., 1., 1.]), zero_point=tensor([0, 0, 0])
      (activation_post_process): MinMaxObserver(averaging_constant=0.01)
    )
  )
)

Each layer out qconfig:
+------------------+----------------------------------------------------+--------------------+-----------------+-------------------+-----------------------------------------+
| Module Name      | Module Type                                        | Input dtype        | out dtype       | ch_axis           | observer                                |
|------------------+----------------------------------------------------+--------------------+-----------------+-------------------+-----------------------------------------|
| quantx           | horizon_plugin_pytorch.nn.qat.stubs.QuantStub      | [torch.float32]    | ['qint8']       | -1                | MinMaxObserver(averaging_constant=0.01) |
| quanty           | horizon_plugin_pytorch.nn.qat.stubs.QuantStub      | [torch.float32]    | ['qint8']       | -1                | MinMaxObserver(averaging_constant=0.01) |
| _generated_add_0 | horizon_plugin_pytorch.nn.qat.conv2d.ConvAddReLU2d | ['qint8', 'qint8'] | [torch.float32] | activation = None |                                         |
| dequant          | horizon_plugin_pytorch.nn.qat.stubs.DeQuantStub    | [torch.float32]    | [torch.float32] | qconfig = None    |                                         |
+------------------+----------------------------------------------------+--------------------+-----------------+-------------------+-----------------------------------------+

可以看到，一组 Module 算子被融合为一个 ConvAddReLU2d，并替换了原始模型中的 _generated_add_0。

因此在进行量化配置方式迁移时，会出现模型打印结果和检查结果变化的情况，需要注意不能机械地进行结果对齐。

可以融合的算子

目前支持的可融合的算子组合见以下函数定义：

import operator
import torch
from torch import nn
from horizon_plugin_pytorch import nn as horizon_nn


def register_fusion_patterns():
    convs = (
        nn.Conv2d,
        nn.ConvTranspose2d,
        nn.Conv3d,
        nn.Linear,
    )
    bns = (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d, nn.SyncBatchNorm)
    adds = (
        nn.quantized.FloatFunctional.add,
        horizon_nn.quantized.FloatFunctional.add,
        torch.add,
        operator.add,  # 即代码中使用的加号
    )
    relus = (nn.ReLU, nn.ReLU6, nn.functional.relu, nn.functional.relu6)

    for conv in convs:
        for bn in bns:
            for add in adds:
                for relu in relus:
                    # conv bn
                    register_fusion_pattern((bn, conv))(ConvBNAddReLUFusion)

                    # conv relu
                    register_fusion_pattern((relu, conv))(ConvBNAddReLUFusion)

                    # conv add
                    register_fusion_pattern((add, conv, MatchAllNode))(
                        ConvBNAddReLUFusion
                    )  # conv 的输出作为 add 的第一个输入
                    register_fusion_pattern((add, MatchAllNode, conv))(
                        ConvBNAddedReLUFusion
                    )  # conv 的输出作为 add 的第二个输入

                    # conv bn relu
                    register_fusion_pattern((relu, (bn, conv)))(
                        ConvBNAddReLUFusion
                    )

                    # conv bn add
                    register_fusion_pattern((add, (bn, conv), MatchAllNode))(
                        ConvBNAddReLUFusion
                    )
                    register_fusion_pattern((add, MatchAllNode, (bn, conv)))(
                        ConvBNAddedReLUFusion
                    )

                    # conv add relu
                    register_fusion_pattern((relu, (add, conv, MatchAllNode)))(
                        ConvBNAddReLUFusion
                    )
                    register_fusion_pattern((relu, (add, MatchAllNode, conv)))(
                        ConvBNAddedReLUFusion
                    )

                    # conv bn add relu
                    register_fusion_pattern(
                        (relu, (add, (bn, conv), MatchAllNode))
                    )(ConvBNAddReLUFusion)
                    register_fusion_pattern(
                        (relu, (add, MatchAllNode, (bn, conv)))
                    )(ConvBNAddedReLUFusion)

#算子融合

#吸收 BN

#融合 Add、ReLU(6)

#实现原理

#融合 Conv+BN

#融合 Conv+Add/ReLU

#版本迁移说明

#可以融合的算子