高效模型量化工具包-AIMET（七）

章节5

量化感知训练

5.1 介绍

  前几章介绍的后训练量化技术是我们量化工具包中的首选工具。它们非常有效且易于实施，因为它们不需要使用标记数据对网络进行重新训练。然而，对于某些模型或使用情况，这些技术可能不足够。后训练技术可能无法减轻低位量化所带来的显著量化误差。在这些情况下，我们采用量化感知训练（QAT）。
  QAT在训练过程中对量化噪声进行建模，并允许模型找到比后训练量化更好的解决方案。然而，更高的准确性也伴随着神经网络训练的常规成本，即更长的训练时间、需要标记数据和超参数搜索。在本章中，我们简要探讨了在具有模拟量化的网络中如何进行反向传播，并提供了一个有效训练QAT模型的标准流程。

  在第3.1节中，我们看到了如何在深度学习框架中使用浮点数模拟量化。然而，如果我们看一下图3.1的计算图，要训练这样一个网络，我们需要通过模拟的量化器块进行反向传播。这带来了一个问题，因为方程（2.4）中的最近舍入操作的梯度在任何地方都是零或未定义的，这使得基于梯度的训练变得不可能。解决这个问题的方法是使用直通估计器来近似梯度，将舍入操作的梯度近似为直通估计器。使用这个梯度定义，我们现在可以通过量化块进行反向传播。图5.1显示了用于量化感知训练中的前向和后向传递的简单计算图。前向传递与图3.1中的相同，但在后向传递中，由于STE假设，我们实际上跳过了量化器块。权重和激活的量化范围可以使用最小-最大范围在每次迭代中进行更新，或者也可以像LSQ中那样学习范围。

5.2 QAT流程

在本节中，我们介绍了QAT的最佳实践流程。我们在图5.2中展示了推荐的流程。这个流程在各种计算机视觉和自然语言处理模型和任务上都能产生良好的QAT结果，并且可以看作是实现低位量化性能的首选工具。如前几节所讨论的，我们总是从预训练的FP32模型开始，并按照一些PTQ步骤进行，以实现更快的收敛和更高的准确性。

跨层均衡： 与PTQ类似，首先我们对全精度模型应用CLE。如第4节所述，对于权重分布不平衡的模型（如MobileNet架构），这一步是必要的。对于其他网络或通道量化的情况，这一步是可选的。
添加量化器： 接下来，我们根据第3.1节的描述选择量化器并添加量化操作到网络中。量化器的选择可能取决于特定的目标硬件，对于常见的AI加速器，我们建议使用对称量化器用于权重，使用非对称量化器用于激活（参见第2.3节）。如果硬件/软件堆栈支持，那么对于权重，最好使用通道量化。在这个阶段，我们还要确保我们对批归一化的模拟是正确的，如第3.2节所讨论的。AIMET会找到这些插入的量化模拟操作的量化参数，如比例/偏移量。为此，AIMET需要用户提供一个回调方法，通过模型传递一些代表性的数据样本。这些样本可以来自训练或校准数据集。通常，样本数量在1000左右对于AIMET来说已经足够找到最佳的量化参数（视觉任务）。
范围设置： 在训练之前，我们必须初始化所有的量化参数。更好的初始化将有助于更快的训练，并可能提高最终的准确性，尽管改进往往很小。一般来说，我们建议使用基于层的SQNR准则来设置所有的量化参数。在通道量化的特定情况下，使用最小-最大设置有时可能更有利。
可学习的量化参数： 可选地，量化器参数可以设为可学习的。直接学习量化参数，而不是在每个时期更新它们，尤其是在处理低位量化时，会带来更高的性能。
训练： 从AIMET返回量化模拟模型（参见第3.1节），我们可以将其作为原始模型在训练流程中的替代品来执行量化感知训练。
导出： 在微调之后，用户可以使用第3.3节中详细介绍的API导出优化的量化模型。

代码示例
API API下面是使用AIMET量化模拟进行QAT的示例。

#代码块5.1 量化感知训练API示例

import torch
from aimet_torch.examples import mnist_torch_model
# Quantization related import
from aimet_torch.quantsim import QuantizationSimModel
model = mnist_torch_model.Net().to(torch.device(’cuda’))
# create Quantization Simulation model
sim = QuantizationSimModel(model,
	dummy_input=torch.rand(1, 1, 28, 28),
	default_output_bw=8,
	default_param_bw=8)
# Quantize the untrained MNIST model
sim.compute_encodings(
	forward_pass_callback=send_samples,
	forward_pass_callback_args=5)
# Fine-tune the model’s parameter using training
trainer_function(model=sim.model, epochs=1,
	num_batches=100, use_cuda=True)
# Export the model and corresponding quantization encodings
sim.export(path=’./’, filename_prefix=’quantized_mnist',
	dummy_input=torch.rand(1, 1, 28, 28))

图5.2 量化感知训练流程。（蓝色方框表示步骤，绿松石色方框表示推荐选择，灰色方框表示可选步骤）

使用注意事项：我们提供了改进量化模型准确性和收敛时间的一般指导方针，用于量化感知训练：
初始化： 从一个良好初始化的量化模型开始。在训练之前应用后训练量化技术，如CLE。
超参数：
  -训练轮数：通常原始轮数的10-20%足以实现收敛。例如：对于典型的视觉模型，我们观察到15-20轮可以得到良好的结果。
  -学习率：与FP32模型在收敛时的最终学习率相当（或高一个数量级）。AIMET的结果是使用1e-6的学习率。
  -学习率调度：每5-10轮将学习率除以10。
  -优化器配置：当使用可学习的量化器时，我们建议在量化参数上使用Adam/RMS Prop优化器，并遵循其他权重微调的建议。

5.2.1 批归一化折叠和QAT

  在第3.2节中，我们介绍了批归一化折叠，它将缩放和加法操作合并到一个线性层中，以实现更高效的推理。在量化感知训练中，我们希望尽可能地模拟推理行为，因此我们必须考虑批归一化折叠在训练过程中的影响。需要注意的是，在一些QAT文献中，忽略了批归一化折叠的影响，并且在量化操作中不考虑批归一化。对于按通道量化来说，这并不重要，因为在推理过程中，批归一化层应用了固定的按通道重新缩放，可以折叠到按通道量化的比例因子中。然而，对于按张量量化来说，如果我们只在部署时将批归一化折叠到权重张量中，可能会导致显著的准确性下降。这是因为网络已经被训练来适应不同的量化噪声。
  在QAT中，一种简单但有效的建模批归一化折叠的方法是在训练之前将批归一化的缩放和偏移静态地折叠到线性层的权重和偏置中。这相当于对权重进行重新参数化，并且有效地将批归一化操作完全移除了网络。AIMET使用静态批归一化折叠方法。

5.3 结果

表5.1  AIMET对ImageNet进行量化感知训练的结果（top-1准确率）。PTQ和QAT都使用W8/A8量化配置。AIMET QAT是在PTQ初始化的基础上进行的。

表5.2  使用AIMET量化感知训练的DeepSpeech2双向LSTM的词错误率（WER）结果（数值越低越好）。AIMET QAT使用W8/A8量化配置和PTQ初始化

   使用QAT流程，我们对用于PTQ的模型进行了量化和评估（参见第4.7节）。我们的结果在表5.1中展示了两种不同的模型架构——MobileNetV2和ResNet50。
   AIMET QAT缩小了与FP32准确性之间的差距（MobileNetV2的<0.5%的差异）。
   此外，AIMET还支持对循环神经网络（如RNN、LSTM和GRU）进行量化模拟和量化感知训练。使用AIMET的QAT，可以将具有双向LSTM的DeepSpeech2模型量化为8位精度，并且准确性下降最小，如表5.2所示。