5.6.tensorRT基础(2)-学习编译int8模型，对模型进行int8量化

前言

杜老师推出的 tensorRT从零起步高性能部署 课程，之前有看过一遍，但是没有做笔记，很多东西也忘了。这次重新撸一遍，顺便记记笔记。

本次课程学习 tensorRT 基础-学习编译 int8 模型，对模型进行 int8 量化

课程大纲可看下面的思维导图

1. Int8 量化

这节课我们来学习 Int8 量化

Int8 量化是利用 int8 乘法替换 float32 乘法实现性能加速的一种方法

1. 对于常规的模型有：y = kx + b，此时 x、k、b 都是 float32，对于 kx 的计算使用 float32 的乘法

2. 而对于 int8 模型有：y = tofp32(toint8(k) * toint8(x)) +b，其中 int8 * int8 结果为 int16

3. 因此 int8 模型解决的问题是如何将 float32 合理的转换为 int8，使得精度损失最小（比如 KL 散度，重新训练等等）

4. 也因此，经过 int8 量化的精度会受到影响

5. 官方给出的参考值，分类器影响比较小，一般掉 1~2 个点，检测器可能有 3~5 个点左右的降低，这个只是一个参考值，实际情况还是要去测试才知道

在 tensorRT 中如果想要实现 Int8 量化，主要是实现以下 4 个步骤：

1. 配置 setFlag nvinfer1::BuilderFlag::kINT8

2. 实现 Int8EntropyCalibrator 类并继承自 IInt8EntropyCalibrator2

• 类似于 Logger 类，你自己实现覆盖其中的一些函数，通过调用这些函数实现交互和通信

3. 实例化 Int8EntropyCalibrator 并且设置到 config.setInt8Calibrator

4. Int8EntropyCalibrator 的作用是读取并预处理图像数据作为输入

• 标定过程的理解：对于输入图像 A，使用 FP32 推理后得到 P1 再用 INT8 推理得到 P2，调整 int8 权重使得 P1 与 P2 足够的接近
• 因此标定时需要使用一些图像，正常发布时，使用 100 张图左右即可

我们回到案例，首先 gen-onnx.py 提供了一个 ImageNet 中标准分类器的模型，内容如下：

import torch
import torchvision
import cv2
import numpy as np


class Classifier(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.backbone = torchvision.models.resnet18(pretrained=True)
        
    def forward(self, x):
        feature     = self.backbone(x)
        probability = torch.softmax(feature, dim=1)
        return probability
        

imagenet_mean = [0.485, 0.456, 0.406]
imagenet_std  = [0.229, 0.224, 0.225]

image = cv2.imread("workspace/kej.jpg")
image = cv2.resize(image, (224, 224))            # resize
image = image[..., ::-1]                         # BGR -> RGB
image = image / 255.0
image = (image - imagenet_mean) / imagenet_std   # normalize
image = image.astype(np.float32)                 # float64 -> float32
image = image.transpose(2, 0, 1)                 # HWC -> CHW
image = np.ascontiguousarray(image)              # contiguous array memory
image = image[None, ...]                         # CHW -> 1CHW
image = torch.from_numpy(image)                  # numpy -> torch
model = Classifier().eval()

with torch.no_grad():
    probability   = model(image)
    
predict_class = probability.argmax(dim=1).item()
confidence    = probability[0, predict_class]

labels = open("workspace/labels.imagenet.txt").readlines()
labels = [item.strip() for item in labels]

print(f"Predict: {predict_class}, {confidence}, {labels[predict_class]}")

dummy = torch.zeros(1, 3, 224, 224)
torch.onnx.export(
    model, (dummy,), "workspace/classifier.onnx", 
    input_names=["image"], 
    output_names=["prob"], 
    dynamic_axes={"image": {0: "batch"}, "prob": {0: "batch"}},
    opset_version=11
)

在这个案例中我们使用了 resnet18 这个 backbone，我们对其重新封装了一次，为了在导出 onnx 的时候加上 softmax 这个节点，把后处理部分添加到 onnx 中，那接下来就是对 ImageNet 中图像进行预处理的标准流程，这个过程其实是可以使用 warpAffine 在 CUDA 上一个核函数全部实现，性能会得到极大的提升

之后我们对分类模型执行了推理，同时导出了对应的 onnx，遵循动态 shape 只保证 batch 维度的动态，没有宽高的动态，执行的效果如下所示：

可以看到模型的推理打印出来了，同时 onnx 模型也正确导出了，接下来我们去看看 C++ 里面是如何做的，二话不说先去执行下 make run

可以看到正常编译了 int8 模型并执行了推理，说明整个过程起码是没问题的，相比于 pytorch 推理的结果其置信度还是略微有所下降，那 int8 的速度比 FP32 的速度快上不少，由于我们标定的是这个图，推理也是这个图，精度没什么变化，那实际情况，特别是检测器而言，精度下降还是比较严重的

在 build_model 函数中，我们使用了 make_nvshared 来自动管理内存，避免内存泄漏，避免需要手动 destroy，其它部分的内容和之前的都一样，不同的是下面这部分：

config->setFlag(nvinfer1::BuilderFlag::kINT8);

auto preprocess = [](
    int current, int count, const std::vector<std::string>& files, 
    nvinfer1::Dims dims, float* ptensor
){
    printf("Preprocess %d / %d\n", count, current);

    // 标定所采用的数据预处理必须与推理时一样
    int width = dims.d[3];
    int height = dims.d[2];
    float mean[] = {0.406, 0.456, 0.485};
    float std[]  = {0.225, 0.224, 0.229};

    for(int i = 0; i < files.size(); ++i){

        auto image = cv::imread(files[i]);
        cv::resize(image, image, cv::Size(width, height));
        int image_area = width * height;
        unsigned char* pimage = image.data;
        float* phost_b = ptensor + image_area * 0;
        float* phost_g = ptensor + image_area * 1;
        float* phost_r = ptensor + image_area * 2;
        for(int i = 0; i < image_area; ++i, pimage += 3){
            // 注意这里的顺序rgb调换了
            *phost_r++ = (pimage[0] / 255.0f - mean[0]) / std[0];
            *phost_g++ = (pimage[1] / 255.0f - mean[1]) / std[1];
            *phost_b++ = (pimage[2] / 255.0f - mean[2]) / std[2];
        }
        ptensor += image_area * 3;
    }
};

// 配置int8标定数据读取工具
shared_ptr<Int8EntropyCalibrator> calib(new Int8EntropyCalibrator(
    {"kej.jpg"}, input_dims, preprocess
));
config->setInt8Calibrator(calib.get());

首先 setFlag 指定为 kINT8，然后构建了一个 lambda 函数 preprocess 用于预处理，接下来我们实例化了一个 Int8EntropyCalibrator 类

Int8EntropyCalibrator 类主要关注：

1. getBatchSize，告诉引擎，这次标定的 batch 是多少

2. getBatch，告诉引擎，这次标定的输入数据是什么，把指针赋值给 bindings 即可，返回 false 表示没有数据了

3. readCalibrationCache，若从缓存文件加载标定信息，则可避免读取文件和预处理，若该函数返回空指针则表示没有缓存，程序会重新通过 getBatch 重新计算，这个特性基本上没什么用

4. writeCalibrationCache，当标定结束后，会调用该函数，我们可以存储标定后的缓存结果，多次标定可以使用该缓存实现加速

最后 setInt8Calibrator 把标定的数据信息塞进去就完成了，那其他的流程就和之前的一模一样，然后在模型序列化后把 calibrator 的数据缓存下来，这个就是 build model 中 int8 量化所做的事情，那推理其实和 fp32 没有任何区别。

那再实际的使用过程中其实 fp16 用得更多一些，int8 可能用的比较少，但是现在来看的话，由于 tensorRT-8 支持显性量化，如果再加上一些优化策略，调点应该不会太严重

2. 补充知识

2.1 知识点

关于 int8 量化的知识点有：(from 杜老师)

1. 对于 int8 需要配置 setFlag nvinfer1::BuilderFlag::kINT8，并且配置 setInt8Calibrator

2. 对于 Int8EntropyCalibrator 则需要继承自 IInt8EntropyCalibrator2

3. Int8EntropyCalibrator 的作用是读取并预处理图像数据作为输入

• 标定的原理，是通过输入标定图像 I 使用参数 WInt8 推理得到输出结果 PInt8，然后不断调整 WInt8，使得输出 PInt8 与 PFloat32 越接近越好
• 因此标定时通常需要使用一些图像，正常发布时，一般使用 100 张图左右即可

4. 常用的 Calibrator

• Int8EntropyCalibrator2

熵校准选择张量的比例因子来优化量化张量的信息论内容，通常会抑制分布中的异常值。这是当前推荐的熵校准器。默认情况下，校准发生在图层融合之前。推荐用于基于 CNN 的网络。

• Iint8MinMaxCalibrator

该校准器使用激活分布的整个范围来确定比例因子。它似乎更适合 NLP 任务。默认情况下，校准发生在图层融合之前。推荐用于 NVIDIA BERT 等网络。

5. 计算机中的 float 计算量是非常庞大的，而改成 int8 后，计算量相比可以提升数倍

对于实际操作时，input[float32]，w[int8]，bias[float32]，output[float32]

步骤如下：

• input[int8] = to_int8(input[float32])

• y[int16] = input[int8] * w[int8] 此处乘法会由计算机转换成 int16，保证精度

• output[float32] = to_float32(y[int16]) + bias[float32]

所以整个过程只是为了减少 float32 的乘法数量以实现提速

对于 to_int8 的过程，并不是直接的线性缩放，而是经过 KL 散度计算最合适的截断点（最大、最小值），进而进行缩放，使得权重的分布尽可能小的被改变

可以参照 https://on-demand.gputechconf.com/gtc/2017/presentation/s7310-8-bit-inference-with-tensorrt.pdf

2.2 其它知识

回过头来看看当前的量化，之前有学习过量化课程，关于 int8 量化笔记可参考 TensorRT量化第四课：PTQ与QAT

首先 tensorRT int8 量化拥有两种模式，分别是 implicity 以及 explicitly 量化，即隐式量化和显示量化，即训练后量化和训练中量化，即 PTQ(Post-Training Quantization) 和 QAT(Quantization Aware Training)

那么本节课程提到的 int8 量化想必你也猜到了就是隐式量化，即 PTQ，它在 tensorRT-7 版本之前用的比较多，而 QAT 则在 tensorRT-8 版本后才完全支持，具体就是可以加载带有 QDQ 信息的模型然后生成对应量化版本的 engine，因此目前来说是不是 QAT 才是主流呢

PTQ 量化就是我们本节课程提到的，它不需要训练，只需要提供一些样本图片，然后在已经训练好的模型上进行校准，统计出来需要的每一层 scale 就可以实现量化了，大概流程如下：

• 在准备好的校准数据集上评估预训练模型

• 使用校准数据来校准模型(校准数据可以是训练集的子集)

• 计算网络中权重和激活的动态范围用来算出量化参数(q-params)

• 使用 q-params 量化网络并执行推理

其中校准方法有什么 entropy、minmax 等等，通过这些校准算法进行 PTQ 量化时，tensorRT 会在优化网络的时候尝试 INT8 精度，假设某一层在 INT8 精度下速度优于默认精度 (FP32或者FP16)，则有限使用 INT8

PTQ 缺点是量化中我们无法控制某一层的精度，毕竟是闭源的，而且 tensorRT 是以速度优化为优先的，很可能某一层你想让它跑 INT8 结果却是 FP16。当然 PTQ 优点是流程简单，速度快。

值得一提的是，在 tensorRT_Pro 这个 repo 中的 int8 量化也是使用的隐式量化

QAT 量化是 tensorRT8 的一个新特性，这个特性其实是指 tensorRT 有直接加载 QAT 模型的能力。而 QAT 模型在这里是指包含 QDQ 操作的量化模型，而 QDQ 操作就是指量化和反量化操作。

实际上 QAT 过程和 tensorRT 没有太大关系，tensorRT 只是一个推理框架，实际的训练中量化操作一般都是在训练框架中去做，比如我们熟悉的 Pytorch。(当然也不排除之后一些推理框架也会有训练功能，因此同样可以在推理框架中做)

tensorRT-8 可以显式地加载包含有 QAT 量化信息的 ONNX 模型，实现一系列优化后，可以生成 INT8 的 engine。

QAT 量化需要插入 QAT 算子且需要训练进行微调，大概流程如下

• 准备一个预训练模型
• 在模型中添加 QAT 算子
• 微调带有 QAT 算子的模型
• 将微调后模型的量化参数即 q-params 存储下来
• 量化模型执行推理

那关于 QAT 相关的知识这里就不过多介绍了，大家感兴趣的可以自行去了解

总结

本节课程学习了 int8 量化，它主要是通过校准算法去进行优化的，不断调整权重使得 int8 推理结果和 fp32 接近。那本节课程学习的 int8 量化本质是 PTQ 隐式量化，而随着 tensorRT-8 版本的发布后 QAT 量化可能慢慢的会成为主流，毕竟可以自己来控制每一层的精度，可操作性高。