TensorRT量化

发布时间:2024年01月15日

系列文章目录

第一章 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 <opencv2/opencv.hpp>

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 <cuda_runtime.h>

// 假设在设备上分配了一块内存 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加速过程。

文章来源:https://blog.csdn.net/m0_70420861/article/details/135536173
本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。