系列文章目录
第一章 YOLOv5模型训练集标注、训练流程
第二章 YOLOv5模型转ONNX,ONNX转TensorRT Engine
第三章 TensorRT量化
文章目录
- 系列文章目录
- 前言
- 一、量化
- 二、量化在TensorRT中的实现
- 三、预处理(Preprocess)和后处理(Postprocess)
- 总结
前言
学习笔记–恩培老师
一、量化
深度学习量化是将深度学习模型中的参数(例如权重)从浮点数转换成整理或者定点数的过程,这样做可以减少模型的存储和计算成本,从而达到模型压缩和运算加速的目的。
- 加快推理速度:访问一次32位浮点型可以访问4次Int8整型数据;
- 减少存储空间和内存占用:在边缘设备(如嵌入式)上部署更实用。
提升速度的同时,量化也会带来精度损失,为了尽可能减少量化过程中精度的损失,需要使用各种校准方法来降低信息的损失。TensorRT中支持两种INT8整型数据:
- 熵校准
- 最小最大值校准
熵校准是一种动态校准算法,使用KL散度(KL Divergence)来度量推理数据和校准数据之间的分布差异。在熵校准中,校准数据是从实时推理数据中采集,它将INT8精度量化参数看作概率分布,根据推理数据和校准数据的KL散度来更新量化参数。这种方法优点是可以更好地反映实际推理数据的分布。
最小最大值校准使用最小最大值算法来计算量化参数。在最小最大值校准中,需要使用一组代表性的校准数据来生成量化参数,首先将推理中的数据进行统计,计算数据的最小值和最大值,然后根据这些值来计算量化参数。
一般选择500-1000张数据用于量化。
二、量化在TensorRT中的实现
在 TensorRT 中,可以通过实现 IInt8EntropyCalibrator2 接口或者 IInt8MinMaxCalibrator 接口来执行熵校准或最小最大值校准。
- IInt8EntropyCalibrator2接口是用于执行熵校准的。需要实现该接口,并提供一个数据生成器,该生成器生成用于校准的数据。在校准过程中,TensorRT将分析每个张量的分布并选择合适的量化参数。你可以在校准过程中自定义数据生成器,例如从文件加载数据集并进行预处理。
- IInt8MinMaxCalibrator接口是用于执行最小最大值校准的。需要实现该接口,并提供一个数据生成器,该生成器生成用于校准的数据。在校准过程中,TensorRT 将计算每个张量的最小和最大值,并使用它们作为量化参数。你可以在校准过程中自定义数据生成器,例如从文件加载数据集并进行预处理。
以下是一个使用 IInt8EntropyCalibrator2 接口的示例:
import tensorrt as trtclass EntropyCalibrator(trt.IInt8EntropyCalibrator2):def __init__(self, data_dir, batch_size, input_shape):# 初始化数据生成器self.data_dir = data_dirself.batch_size = batch_sizeself.input_shape = input_shapedef get_batch_size(self):return self.batch_sizedef 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预处理步骤如下
- letterbox: 即保持原图比例(图像直接resize到输入大小效果不好),将图片放在一个正方形画布中,多余部分用黑色填充。
- Normalization(归一化):将像素值缩放到[0,1]间;
- 颜色通道顺序调整:BGR2RGB
- 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;
}
后处理
- 在设备上分配主机内存,用于存储复制后的数据。
- 使用cudaMemcpy函数将设备上的结果复制到分配的主机内存中。
- 释放在设备上分配的内存。
#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加速过程。