AIMET API 文档（1）

1 适用于 PyTorch 的 AIMET API
- 1.1 AIMET PyTorch 量化 API
- - 1.1.1 PyTorch 模型指南
  - 1.1.2 架构检查器 API

1 适用于 PyTorch 的 AIMET API

1.1 AIMET PyTorch 量化 API

为了充分利用 AIMET 量化功能，鼓励用户在定义 PyTorch 模型时遵循几条准则。 AIMET 提供的 API 可以自动执行一些模型定义更改，并检查 AIMET 量化功能是否可以应用于 PyTorch 模型。

在使用任何 AIMET 量化功能之前，用户应首先调用模型准备器 API。

模型指南：定义 PyTorch 模型的指南
架构检查器 API：：允许用户检查模型的性能问题
模型准备器 API：允许用户自动更改模型定义
模型验证器 API：允许用户检查 AIMET 量化功能是否可以应用于 PyTorch 模型

PyTorch 模型的 AIMET 量化提供以下功能。

定量分析器 API：分析模型并指出量化的敏感层
量化模拟 API：能够在量化硬件上模拟推理和训练
自适应舍入 API：训练后量化技术，用于优化权重张量的舍入
跨层均衡 API：用于均衡层参数的训练后量化技术
偏差校正 API：训练后量化技术，用于校正由于量化噪声而导致的层输出偏移
AutoQuant API：集成了AIMET提供的训练后量化技术的统一API
BN Re-estimation API：重新估计 BN 层统计数据并折叠 BN 层的 API API：集成了AIMET提供的训练后量化技术的统一API

如果用户想将Multi-GPU与CLE或QAT结合使用，可以参考：

多 GPU 指南：将 PyTorch DataParallel API 与 AIMET 功能结合使用的指南

1.1.1 PyTorch 模型指南

为了充分利用 AIMET 功能，鼓励用户在定义 PyTorch 模型时遵循几条准则。

模型应支持转换为 onnx

模型定义应支持转换为onnx，用户可以检查模型对onnx转换的兼容性，如下所示：

...
model = Model()
torch.onnx.export(model, <dummy_input>, <onnx_file_name>):

模型应该是 jit 可追踪的

模型定义应该是jit可追踪的，用户可以检查jit追踪模型的兼容性，如下所示：

...
model = Model()
torch.jit.trace(model, <dummy_input>):

将层定义为模块，而不是使用 torch.nn.function 等效项

当使用激活函数和其他无状态层时，PyTorch 将允许用户

将层定义为模块（在构造函数中实例化并在前向传递中使用），或者
纯粹在前向传播中使用 torch.nn.function 等价物

对于AIMET量化仿真模型添加仿真节点，AIMET需要前者（层定义为模块）。更改模型定义以使用模块而不是函数在数学上是等效的，并且不需要重新训练模型。

举个例子，如果用户有：

def forward(...):
    ...
    x = torch.nn.functional.relu(x)
    ...

用户应该将他们的模型定义为：

def __init__(self,...):
    ...
    self.relu = torch.nn.ReLU()
    ...

def forward(...):
    ...
    x = self.relu(x)
    ...

在某些情况下，这是不可能的，因为操作只能表示为函数而不是类定义，但应尽可能遵循。

此外，用户还可以使用模型准备器 API 来自动化此操作

避免重用类定义的模块

类定义中定义的模块只能使用一次。如果重用任何模块，请在类定义中定义一个新的相同模块。例如，如果用户有：

def __init__(self,...):
    ...
    self.relu = torch.nn.ReLU()
    ...

def forward(...):
    ...
    x = self.relu(x)
    ...
    x2 = self.relu(x2)
    ...

用户应该将他们的模型定义为：

def __init__(self,...):
    ...
    self.relu = torch.nn.ReLU()
    self.relu2 = torch.nn.ReLU()
    ...

def forward(...):
    ...
    x = self.relu(x)
    ...
    x2 = self.relu2(x2)
    ...

此外，用户还可以使用模型准备器 API 来自动化此操作

仅使用 torch.Tensor 或 torch.Tensors 的元组作为模型/子模块输入和输出

模块应使用张量或张量元组作为输入和输出，以支持模型到 onnx 的转换。例如，如果用户有：

def __init__(self,...):
...
def forward(self, inputs: Dict[str, torch.Tensor]):
    ...
    x = self.conv1(inputs['image_rgb'])
    rgb_output = self.relu1(x)
    ...
    x = self.conv2(inputs['image_bw'])
    bw_output = self.relu2(x)
    ...
    return {
    
     'rgb': rgb_output, 'bw': bw_output }

用户应该将他们的模型定义为：

def __init__(self,...):
...
def forward(self, image_rgb, image_bw):
    ...
    x = self.conv1(image_rgb)
    rgb_output = self.relu1(x)
    ...
    x = self.conv2(image_bw)
    bw_output = self.relu2(x)
    ...
    return rgb_output, bw_output

1.1.2 架构检查器 API

aimet_torch.arch_checker.arch_checker.ArchChecker.check_model_arch(model, dummy_input, result_dir=None)

使用 _node_check_dict 中的检查检查模型中的每个节点。仅记录节点和失败的测试。

参数：

model（Module） - 要进行检查的Torch模型。
dummy_input (Union[Tensor, Tuple]) – 将输入传递给模型。可以是张量或张量元组

返回 arch_checker_report ：{op.dotted_name_op: NodeErrorReportObject }
返回类型：ArchCheckerReport

AIMET PyTorch 架构检查器有助于检查次优模型构造，并提供更新模型以提高性能的潜在选项。架构检查器当前检查以下条件：

用于最佳通道大小的卷积层。
不具备性能的激活函数。
批量归一化层不能折叠。
具有填充的卷积层序列中的中间卷积层。

在本节中，我们将介绍未通过架构检查的模型，并展示如何运行架构检查器。

示例 1：通道不足的模型

我们从以下模型开始，其中包含通道数小于 32 的卷积层。

class ModelWithNotEnoughChannels(torch.nn.Module):
    """ Model that prelu module. Expects input of shape (1, 3, 32, 32) """

    def __init__(self):
        super(ModelWithNotEnoughChannels, self).__init__()
        self.conv1 = torch.nn.Conv2d(3, 31, kernel_size=2, stride=2, padding=2, bias=False)
        self.bn1 = torch.nn.BatchNorm2d(31)

    def forward(self, *inputs):
        x = self.conv1(inputs[0])
        x = self.bn1(x)
        return x

导入架构检查器：

从aimet_torch.arch_checker.arch_checker导入ArchChecker

通过传入模型以及模型输入来运行模型检查器：

def example_check_for_number_of_conv_channels():

    model = ModelWithNotEnoughChannels()
    ArchChecker.check_model_arch(model, dummy_input=torch.rand(1, 3, 32, 32))

模型中的卷积层少了一个通道，将出现以下记录器打印：

Utils - INFO - Graph/Node: ModelWithNotEnoughChannels.conv1: Conv2d(3, 31, kernel_size=(2, 2), stride=(2, 2), padding=(2, 2), bias=False) fails check: {
    
    '_check_conv_channel_32_base', '_check_conv_channel_larger_than_32'}

将生成包含以下内容的 HTML 文件。
HTML 报告内容

图/图层名称	问题	推荐
ModelWithNotEnoughChannels.conv1	该卷积的输入/输出张量的通道大小小于32	尝试将通道调整为 32 的倍数以获得更好的性能
ModelWithNotEnoughChannels.conv1	该卷积的输入/输出张量的通道大小不是32的倍数	尝试将通道调整为 32 的倍数以获得更好的性能。

示例 2：具有非性能激活的模型

我们从以下模型开始，其中包含通道数小于 32 的卷积层。

class ModelWithPrelu(torch.nn.Module):
    """ Model that prelu module. Expects input of shape (1, 3, 32, 32) """

    def __init__(self):
        super(ModelWithPrelu, self).__init__()
        self.conv1 = torch.nn.Conv2d(32, 32, kernel_size=2, stride=2, padding=2, bias=False)
        self.bn1 = torch.nn.BatchNorm2d(32)
        self.prelu1 = torch.nn.PReLU()

    def forward(self, *inputs):
        x = self.conv1(inputs[0])
        x = self.bn1(x)
        x = self.prelu1(x)
        return x

通过传入模型以及模型输入来运行模型检查器：

def example_check_for_non_performant_activations():

    model = ModelWithPrelu()
    ArchChecker.check_model_arch(model, dummy_input=torch.rand(1, 32, 32, 32))

与 ReLU 相比，模型中的 PReLU 层被认为性能不佳，将出现以下记录器打印：

Utils - INFO - Graph/Node: ModelWithPrelu.prelu1: PReLU(num_parameters=1) fails check: {
    
    '_activation_checks'}

示例 3：具有独立批量归一化层的模型

我们从以下模型开始，其中包含通道数小于 32 的卷积层。

class ModelWithNonfoldableBN(torch.nn.Module):
    """ Model that has non-foldable batch norm. """

    def __init__(self):
        super(ModelWithNonfoldableBN, self).__init__()
        self.conv1 = torch.nn.Conv2d(32, 32, kernel_size=2, stride=2, padding=2, bias=False)
        self.avg_pool1 = torch.nn.AvgPool2d(3, padding=1)
        self.bn1 = torch.nn.BatchNorm2d(32)

    def forward(self, *inputs):
        x = self.conv1(inputs[0])
        x = self.avg_pool1(x)
        x = self.bn1(x)
        return x

通过传入模型以及模型输入来运行模型检查器：

def example_check_for_standalone_bn():

    model = ModelWithNonfoldableBN()
    ArchChecker.check_model_arch(model, dummy_input=torch.rand(1, 32, 32, 32))

AveragePool 层防止 BatchNormalization 层与 Convolution 层折叠，将出现以下记录器打印：

Utils - INFO - Graph/Node: ModelWithNonfoldableBN.bn1: BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) fails check: {
    
    '_check_batch_norm_fold'}