应用感知量化训练

Linux Ascend GPU 模型调优 高级

背景

越来越多的应用选择在移动设备或者边缘设备上使用深度学习技术。以手机为例,为了提供人性化和智能的服务,现在操作系统和应用都开始集成深度学习功能。而使用该功能,涉及训练或者推理,自然包含大量的模型及权重文件。经典的AlexNet,原始权重文件已经超过了200MB,而最近出现的新模型正往结构更复杂、参数更多的方向发展。由于移动设备、边缘设备的硬件资源有限,需要对模型进行精简,而量化(Quantization)技术就是应对该类问题衍生出的技术之一。

概念

量化

量化即以较低的推理精度损失将连续取值(或者大量可能的离散取值)的浮点型模型权重或流经模型的张量数据定点近似(通常为INT8)为有限多个(或较少的)离散值的过程,它是以更少位数的数据类型用于近似表示32位有限范围浮点型数据的过程,而模型的输入输出依然是浮点型。这样的好处是可以减小模型尺寸大小,减少模型内存占用,加快模型推理速度,降低功耗等。

如上所述,与FP32类型相比,FP16、INT8、INT4等低精度数据表达类型所占用空间更小。使用低精度数据表达类型替换高精度数据表达类型,可以大幅降低存储空间和传输时间。而低比特的计算性能也更高,INT8相对比FP32的加速比可达到3倍甚至更高,对于相同的计算,功耗上也有明显优势。

当前业界量化方案主要分为两种:感知量化训练(Quantization Aware Training)和训练后量化(Post-training Quantization)。

伪量化节点

伪量化节点,是指感知量化训练中插入的节点,用以寻找网络数据分布,并反馈损失精度,具体作用如下:

  • 找到网络数据的分布,即找到待量化参数的最大值和最小值;

  • 模拟量化为低比特时的精度损失,把该损失作用到网络模型中,传递给损失函数,让优化器在训练过程中对该损失值进行优化。

感知量化训练

MindSpore的感知量化训练是在训练基础上,使用低精度数据替换高精度数据来简化训练模型的过程。这个过程不可避免引入精度的损失,这时使用伪量化节点来模拟引入的精度损失,并通过反向传播学习,来减少精度损失。对于权值和数据的量化,MindSpore采用了参考文献[1]中的方案。

感知量化训练规格

规格 规格说明
硬件支持 GPU、Ascend AI 910处理器的硬件平台
网络支持 已实现的网络包括LeNet、ResNet50等网络,具体请参见https://gitee.com/mindspore/mindspore/tree/r1.1/model_zoo
算法支持 在MindSpore的伪量化训练中,支持非对称和对称的量化算法。
方案支持 支持4、7和8比特的量化方案。
数据类型支持 Ascend平台支持精度为FP32和FP16的网络进行量化训练,GPU平台支持FP32。

感知量化训练示例

感知量化训练模型与一般训练步骤一致,在定义网络和最后生成模型阶段后,需要进行额外的操作,完整流程如下:

  1. 数据处理加载数据集。

  2. 定义原始非量化网络。

  3. 定义融合网络。在完成定义原始非量化网络后,替换指定的算子,完成融合网络的定义。

  4. 定义优化器和损失函数。

  5. 转化量化网络。基于融合网络,使用转化接口在融合网络中插入伪量化节点,生成量化网络。

  6. 进行量化训练。基于量化网络训练,生成量化模型。

在上面流程中,第3、5、6步是感知量化训练区别普通训练需要额外进行的步骤。

  • 融合网络:使用指定算子(nn.Conv2dBnActnn.DenseBnAct)替换后的网络。

  • 量化网络:融合模型使用转换接口(QuantizationAwareTraining.quantize)插入伪量化节点后得到的网络。

  • 量化模型:量化网络训练后得到的checkpoint格式的模型。

接下来,以LeNet网络为例,展开叙述2、3两个步骤。

你可以在这里找到完整可运行的样例代码:https://gitee.com/mindspore/mindspore/tree/r1.1/model_zoo/official/cv/lenet_quant

定义融合网络

定义融合网络,在定义网络后,替换指定的算子。

  1. 使用nn.Conv2dBnAct算子替换原网络模型中的2个算子nn.Conv2dnn.ReLU

  2. 使用nn.DenseBnAct算子替换原网络模型中的2个算子nn.Densenn.ReLU

无论nn.Densenn.Conv2d算子后面有没有nn.BatchNorm*nn.ReLU,都要按规定使用上述两个算子进行融合替换。

原网络模型LeNet5的定义如下所示:

class LeNet5(nn.Cell):
    """
    Lenet network

    Args:
        num_class (int): Num classes. Default: 10.
        num_channel (int): Num channel. Default: 1.
    Returns:
        Tensor, output tensor
    Examples:
        >>> LeNet(num_class=10, num_channel=1)

    """
    def __init__(self, num_class=10, num_channel=1):
        super(LeNet5, self).__init__()
        self.conv1 = nn.Conv2d(num_channel, 6, 5, pad_mode='valid')
        self.conv2 = nn.Conv2d(6, 16, 5, pad_mode='valid')
        self.fc1 = nn.Dense(16 * 5 * 5, 120, weight_init=Normal(0.02))
        self.fc2 = nn.Dense(120, 84, weight_init=Normal(0.02))
        self.fc3 = nn.Dense(84, num_class, weight_init=Normal(0.02))
        self.relu = nn.ReLU()
        self.max_pool2d = nn.MaxPool2d(kernel_size=2, stride=2)
        self.flatten = nn.Flatten()

    def construct(self, x):
        x = self.max_pool2d(self.relu(self.conv1(x)))
        x = self.max_pool2d(self.relu(self.conv2(x)))
        x = self.flatten(x)
        x = self.relu(self.fc1(x))
        x = self.relu(self.fc2(x))
        x = self.fc3(x)
        return x

替换算子后的融合网络如下:

class LeNet5(nn.Cell):
    def __init__(self, num_class=10):
        super(LeNet5, self).__init__()
        self.num_class = num_class

        self.conv1 = nn.Conv2dBnAct(1, 6, kernel_size=5, activation='relu')
        self.conv2 = nn.Conv2dBnAct(6, 16, kernel_size=5, activation='relu')

        self.fc1 = nn.DenseBnAct(16 * 5 * 5, 120, activation='relu')
        self.fc2 = nn.DenseBnAct(120, 84, activation='relu')
        self.fc3 = nn.DenseBnAct(84, self.num_class)
        self.max_pool2d = nn.MaxPool2d(kernel_size=2, stride=2)

    def construct(self, x):
        x = self.max_pool2d(self.conv1(x))
        x = self.max_pool2d(self.conv2(x))
        x = self.flattern(x)
        x = self.fc1(x)
        x = self.fc2(x)
        x = self.fc3(x)
        return x

转化为量化网络

使用QuantizationAwareTraining.quantize接口自动在融合模型中插入伪量化节点,将融合模型转化为量化网络。

from mindspore.compression.quant import QuantizationAwareTraining

quantizer = QuantizationAwareTraining(quant_delay=900,
                                      bn_fold=False,
                                      per_channel=[True, False],
                                      symmetric=[True, False])
net = quantizer.quantize(network)

重训和推理

导入模型重新训练

上面介绍了从零开始进行感知量化训练。更常见情况是已有一个模型文件,希望生成量化模型,这时已有正常网络模型训练得到的模型文件及训练脚本,进行感知量化训练。这里使用checkpoint文件重新训练的功能,详细步骤为:

  1. 数据处理加载数据集。

  2. 定义原始非量化网络。

  3. 训练原始网络生成非量化模型。

  4. 定义融合网络。

  5. 定义优化器和损失函数。

  6. 基于融合网络转化生成量化网络。

  7. 加载模型文件重训。加载已有非量化模型文件,基于量化网络重新训练生成量化模型。详细模型重载训练,请参见https://www.mindspore.cn/tutorial/training/zh-CN/r1.1/use/load_model_for_inference_and_transfer.html

进行推理

使用量化模型进行推理,与普通模型推理一致,分为直接checkpoint文件推理及转化为通用模型格式(AIR、MindIR等)进行推理。

  • 使用感知量化训练后得到的checkpoint文件进行推理:

    1. 加载量化模型。

    2. 推理。

  • 转化为ONNX等通用格式进行推理(暂不支持,开发完善后补充)。

参考文献

[1] Jacob B, Kligys S, Chen B, et al. Quantization and training of neural networks for efficient integer-arithmetic-only inference[C]//Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition. 2018: 2704-2713.

[2] Krishnamoorthi R. Quantizing deep convolutional networks for efficient inference: A whitepaper[J]. arXiv preprint arXiv:1806.08342, 2018.