Eager 模式

horizon_plugin_pytorch 目前支持采用 eager 模式进行量化,但是,我们已经不推荐再使用该模式。 Eager 模式的整体流程同样参考了 PyTorch 官方的量化接口和思路,因此,建议您先阅读 PyTorch 官方文档中 Eager 模式相关部分。

与 fx 模式的区别

在 horizon_plugin_pytorch 中使用 eager 模式,和 fx 模式的主要区别在于:

  • eager 模式仅支持 module 形式的算子。在执行量化流程前,您需要手动将浮点模型中的函数形式的算子替换为 PyTorch 中 Module 类型的算子或者是 horizon_plugin_pytorch 中定义的专有算子,包括但不限于:
原始的浮点算子需要替换的算子
torch.nn.functional.relutorch.nn.ReLU()
a + b /
torch.add
horizon.nn.quantized.FloatFunctional().add
Tensor.exphorizon.nn.Exp()
torch.nn.functional.interpolatehorizon.nn.Interpolate()
  • 您必须手动定义需要融合的算子,并在执行量化流程前显式调用融合函数执行算子融合,调用时也需指定使用 horizon_plugin_pytorch 中提供的 fuser_func。如下所示:
import torch
from torch import nn
import horizon_plugin_pytorch as horizon


class ConvBNReLU(nn.Sequential):
    def __init__(self, in_channels, out_channels, kernel_size):
        super(ConvBNReLU, self).__init__(
            nn.Conv2d(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=kernel_size
            ),
            nn.BatchNorm2d(num_features=out_channels),
            nn.ReLU()
        )

    # 指定可以 fuse 的算子
    def fuse_model(self):
        torch.quantization.fuse_modules(
            self,
            ['0', '1', '2'],
            inplace=True,
            # 指定 horizon_plugin_pytorch 中提供的 fuse 函数
            fuser_func=horizon.quantization.fuse_known_modules,
        )

float_model = ConvBNReLU(1, 1, 1)
# 需要显式调用 fuse 函数
float_model.fuse_model()

print(float_model)
# ConvBNReLU(
#   (0): ConvReLU2d(
#     (0): Conv2d(1, 1, kernel_size=(1, 1), stride=(1, 1))
#     (1): ReLU()
#   )
#   (1): Identity()
#   (2): Identity()
# )