TensorRT量化

系列文章目录

第一章 YOLOv5模型训练集标注、训练流程
第二章 YOLOv5模型转ONNX,ONNX转TensorRT Engine
第三章 TensorRT量化

---

---

前言

学习笔记–恩培老师

---

一、量化

深度学习量化是将深度学习模型中的参数（例如权重）从浮点数转换成整理或者定点数的过程，这样做可以减少模型的存储和计算成本，从而达到模型压缩和运算加速的目的。

• 加快推理速度:访问一次32位浮点型可以访问4次Int8整型数据；
• 减少存储空间和内存占用：在边缘设备(如嵌入式）上部署更实用。

提升速度的同时，量化也会带来精度损失，为了尽可能减少量化过程中精度的损失，需要使用各种校准方法来降低信息的损失。TensorRT中支持两种INT8整型数据：

• 熵校准
• 最小最大值校准

熵校准是一种动态校准算法，使用KL散度(KL Divergence)来度量推理数据和校准数据之间的分布差异。在熵校准中，校准数据是从实时推理数据中采集，它将INT8精度量化参数看作概率分布，根据推理数据和校准数据的KL散度来更新量化参数。这种方法优点是可以更好地反映实际推理数据的分布。

最小最大值校准使用最小最大值算法来计算量化参数。在最小最大值校准中，需要使用一组代表性的校准数据来生成量化参数，首先将推理中的数据进行统计，计算数据的最小值和最大值，然后根据这些值来计算量化参数。

一般选择500-1000张数据用于量化。

二、量化在TensorRT中的实现

在 TensorRT 中，可以通过实现 IInt8EntropyCalibrator2 接口或者 IInt8MinMaxCalibrator 接口来执行熵校准或最小最大值校准。

• IInt8EntropyCalibrator2接口是用于执行熵校准的。需要实现该接口，并提供一个数据生成器，该生成器生成用于校准的数据。在校准过程中，TensorRT将分析每个张量的分布并选择合适的量化参数。你可以在校准过程中自定义数据生成器，例如从文件加载数据集并进行预处理。
• IInt8MinMaxCalibrator接口是用于执行最小最大值校准的。需要实现该接口，并提供一个数据生成器，该生成器生成用于校准的数据。在校准过程中，TensorRT 将计算每个张量的最小和最大值，并使用它们作为量化参数。你可以在校准过程中自定义数据生成器，例如从文件加载数据集并进行预处理。

以下是一个使用 IInt8EntropyCalibrator2 接口的示例：

import tensorrt as trt

class EntropyCalibrator(trt.IInt8EntropyCalibrator2):
    def __init__(self, data_dir, batch_size, input_shape):
        # 初始化数据生成器
        self.data_dir = data_dir
        self.batch_size = batch_size
        self.input_shape = input_shape

    def get_batch_size(self):
        return self.batch_size

    def get_batch(self, names):
        # 从数据生成器中获取一个批次的数据
        # 返回一个包含每个输入名称和对应数据的字典
        batch_data = ...
        return batch_data

# 创建 TensorRT builder 和配置
builder = trt.Builder(...)
config = builder.create_builder_config()

# 设置 Int8 校准器
calibrator = EntropyCalibrator(data_dir, batch_size, input_shape)
config.int8_calibrator = calibrator

# 构建和优化 TensorRT 引擎
network = builder.create_network()
engine = builder.build_engine(network, config)

使用 IInt8MinMaxCalibrator 接口也类似，只需实现不同的接口方法。

运行代码

#将会读取 c3.mp4 文件，并将其每一帧保存为以 sample0001.png、sample0002.png、
#sample0003.png等命名的 PNG 图像文件。
ffmpeg -i c3.mp4 sample%04d.png 

#ls 命令获取当前目录下所有的 PNG 图像文件，并使用 shuf 命令随机选择其中的 200 个文件，
#并将结果保存到 filelist.txt 文件中。
ls *.png | shuf -n 200 > filelist.txt

#将Build后的参数分布是onnx,校准目录(用于拼接完整图片路径),文件列表路径
./build/build weights/yolov5s_person.onnx ./media/ ./media/filelist.txt

三、预处理（Preprocess)和后处理(Postprocess)

YOLOv5预处理步骤如下

1. letterbox: 即保持原图比例(图像直接resize到输入大小效果不好),将图片放在一个正方形画布中，多余部分用黑色填充。
2. Normalization（归一化）：将像素值缩放到[0,1]间；
3. 颜色通道顺序调整：BGR2RGB
4. NHWC转为NCHW

letterbox的实现，可以使用opencv的cv::warpAffine

#include 

using namespace cv;

Mat letterbox(Mat& image, int width, int height)
{
    // 获取原始图像的宽度和高度
    int img_width = image.cols;
    int img_height = image.rows;

    // 计算需要添加的填充大小
    float scale = std::min(static_cast<float>(width) / img_width, static_cast<float>(height) / img_height);
    int new_width = static_cast<int>(img_width * scale);
    int new_height = static_cast<int>(img_height * scale);
    int dx = (width - new_width) / 2;
    int dy = (height - new_height) / 2;

    // 创建变换矩阵
    Mat M = Mat::eye(2, 3, CV_32F);
    M.at<float>(0, 0) = scale;
    M.at<float>(1, 1) = scale;
    M.at<float>(0, 2) = dx;
    M.at<float>(1, 2) = dy;

    // 应用变换矩阵
    Mat letterboxed_image;
    warpAffine(image, letterboxed_image, M, Size(width, height), INTER_LINEAR, BORDER_CONSTANT, Scalar(0, 0, 0));

    return letterboxed_image;
}

int main()
{
    // 加载图像
    Mat image = imread("input.jpg");

    // 进行 Letterbox 处理
    Mat letterboxed_image = letterbox(image, 800, 800);

    // 显示调整后的图像并保存
    imshow("Letterboxed Image", letterboxed_image);
    imwrite("output.jpg", letterboxed_image);

    waitKey(0);
    return 0;
}

后处理

1. 在设备上分配主机内存，用于存储复制后的数据。
2. 使用cudaMemcpy函数将设备上的结果复制到分配的主机内存中。
3. 释放在设备上分配的内存。

#include 

// 假设在设备上分配了一块内存 device_output，存储了处理后的结果

// 获取结果的大小和其他信息
size_t output_size = ...;
// 其他信息...

// 在主机上分配内存，用于存储复制后的数据
void* host_output = malloc(output_size);

// 将设备上的结果复制到主机内存中
cudaMemcpy(host_output, device_output, output_size, cudaMemcpyDeviceToHost);

// 进行后续处理，如执行非极大值抑制等

// 释放在主机上分配的内存
free(host_output);

---

总结

接下来会介绍TensorRT结合DeepStream加速过程。