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() # )