YOLOv7-QAT量化部署
目录
-
- 前言
- 一、QAT量化浅析
- 二、YOLOv7模型训练
-
- 1. 项目的克隆和必要的环境依赖
-
- 1.1 项目的克隆
- 1.2 项目代码结构整体介绍
- 1.3 环境安装
- 2. 数据集和预训练权重的准备
-
- 2.1 数据集
- 2.2 预训练权重准备
- 3. 训练模型
-
- 3.1 修改模型配置文件
- 3.2 修改数据配置文件
- 3.3 训练模型
- 3.4 mAP测试
- 三、YOLOv7-QAT准备工作
-
- 1. 项目克隆
- 2. 安装依赖
- 3. 数据集和权重准备
- 4. 配置文件修改
- 四、YOLOv7-QAT微调导出
-
- 1. QAT微调
- 2. QAT模型导出
- 3. INT8模型生成
- 五、YOLOv7-QAT部署
-
- 1. 源码下载
- 2. 环境配置
-
- 2.1 配置CMakeLists.txt
- 2.2 配置Makefile
- 3. QAT模型推理
- 4. QAT模型mAP测试
- 六、讨论
-
- 1. 不同精度模型对比
- 2. PTQ vs. QAT
- 3. YOLOv5-QAT vs. YOLOv7-QAT
- 结语
- 下载链接
- 参考
前言
学习 yolo_deepstream 项目中的 yolov7-qat 量化，本文主要是学习项目中的 YOLOv7 QAT 量化的方法，其他部分如 deepstream 博主并未关注，部署使用的 repo 依旧是 tensorRT_Pro，博主在这里简单的过一遍流程，不涉及任何的原理性分析。
博主为初学者，欢迎交流讨论，若有问题欢迎各位看官批评指正！！！
一、QAT量化浅析
在正式开始之前我们先来回顾下关于 QAT 量化的一些知识，具体可参考：TensorRT量化第四课：PTQ与QAT
TensorRT 有两种量化模式，分别是隐式（implicitly）量化和显式（explicitly）量化。前者在 TRT7 版本之前用得比较多，而后者在 TRT8 版本后才完全支持，具体就是可以加载带有 QDQ 信息的模型然后生成对应量化版本的 engine。
这篇文章主要分享显式量化即 QAT 量化，关于隐式量化即 PTQ 量化可以查看上篇文章：YOLOv7-PTQ量化部署。
QAT（Quantization Aware Training）即训练中量化也叫显式量化。它是 tensorRT8 的一个新特性，这个特性其实是指 tensorRT 有直接加载 QAT 模型的能力。而 QAT 模型在这里是指包含 QDQ 操作的量化模型，而 QDQ 操作就是指量化和反量化操作。
实际上 QAT 过程和 tensorRT 没有太大关系，tensorRT 只是一个推理框架，实际的训练中量化即 QAT 操作一般都是在训练框架中去做的，比如我们熟悉的 Pytorch。（当然也不排除之后一些推理框架也会有训练功能，因此同样可以在推理框架中做）
tensorRT8 可以显式地加载包含有 QAT 量化信息的 ONNX 模型，实现一系列优化后，可以生成 INT8 的 engine。
QAT 量化需要插入 QAT 算子且需要训练进行微调，大概流程如下：
- 准备一个预训练模型
- 在模型中添加 QAT 算子
- 微调带有 QAT 算子的模型
- 将微调后模型的量化参数即 q-params 存储下来
- 量化模型执行推理
带有 QAT 量化信息的模型如下图所示：
从上图中我们可以看到带有 QAT 量化信息的模型中有 QuantizeLinear 和 DequantizeLinear 模块，也就是对应的 QDQ 模块，它包含了该层和该激活值的量化 scale 和 zero-point。什么是 QDQ 呢？QDQ 其实就是 Q（量化）和 DQ（反量化）两个 op，在网络中通常作为模拟量化的 op，如下图所示：
QDQ 模块会参与训练，负责将输入的 FP32 张量量化为 INT8，随后再进行反量化将 INT8 的张量再变为 FP32。值得注意的是，实际网络中训练使用的精度还是 FP32，只不过这个量化算子在训练中可以学习到量化和反量化的尺度信息，这样训练的时候就可以让模型权重和量化参数更好地适应量化过程，量化后地精度也相对更高一些。
QDQ 模块的用途主要体现在两方面：
- 第一个是可以存储量化信息，比如 scale 和 zero_point，这些信息可以放在 Q 和 DQ 操作中
- 第二个是可以当作是显示指定哪一层是量化层，我们可以默认认为包在 QDQ 操作中间的 op 都是 INT8 类型的 op，也就是我们需要量化的 op
因此对比显式量化（即 QAT 量化），tensorRT 的隐式量化（即 PTQ 量化）就没有那么直接，在 tensorRT-8 版本之前我们一般都是借助 tensorRT 的内部量化算法去量化（闭源），在构建 engine 的时候传入图像进行校准，执行的是训练后量化（PTQ）的过程。
而有了 QDQ 信息，tensorRT 在解析模型的时候会根据 QDQ 的位置找到可量化的 op，然后与 QDQ 融合(吸收尺度信息 scale 到 op 中)，融合后的算子就是实打实的 INT8 算子，经过一系列的融合优化后，最终生成量化版的 engine。
OK！关于 QAT 量化我们就简单聊下，让我们开始具体的实现吧！！！
二、YOLOv7模型训练
首先我们需要训练一个 YOLOv7 模型，当然拿官方的预训练权重也行，博主这边为了完整性还是整体走一遍流程，熟悉 YOLOv7 模型训练的看官可以跳过直接到量化部分。
1. 项目的克隆和必要的环境依赖
1.1 项目的克隆
yolov7 的代码是开源的可直接从 github 官网上下载，源码下载地址是 https://github.com/WongKinYiu/yolov7，由于 yolov7 目前就只固定 v0.1 一个版本，而 v0.1 版本并未提供训练的详细说明，故采用主分支进行模型的训练和部署工作。Linux下代码克隆指令如下
git clone https://github.com/WongKinYiu/yolov7.git
也可手动点击下载，点击右上角的绿色的 Code 按键，将代码下载下来。至此整个项目就已经准备好了。也可以点击 here【pwd:yolo】下载博主准备好的代码（注意该代码下载于 2023/10/21 日，若有改动请参考最新）
1.2 项目代码结构整体介绍
将下载后的 yolov7 代码解压，其代码目录如下图所示：
现在来对代码的整体目录做一个介绍：
- |-cfg：存放yolov7不同模型的yaml文件，如yolov7.yaml、yolov7-tiny.yaml等，包括训练和部署时的yolov7模型yaml
- |-data：存放一些超参数的配置文件以及配置训练集和验证集路径的coco.yaml文件，如果需要修改自己的数据集，那么需要修改其中的yaml文件
- |-deploy：针对部署的文件夹
- |-figure：存放yolov7测试的效果图片
- |-inference：存放推理时的图片
- |-models：存放yolov7整体网络模型搭建的py文件
- |-paper：存放yolov7论文
- |-scripts：脚本文件，用于获取coco数据集
- |-tools：该文件夹主要存放一些示例教程，如yolov7关键点检测、yolov7实例分割、yolov7onnx等等
- |-utils：存放工具类函数，包括loss、metrics、plots函数等
- |-
- detect.py：检测代码，包括图像检测、视频流检测等
- export.py：模型导出代码，如onnx导出
- hubconf.py：pytorch扩展模型
- requirements.txt：文本文件，里面包含使用yolov7项目的环境依赖包以及相应的版本号
- test.py：测试代码
- train.py：训练代码
- train_aux.py：训练辅助头代码(不确定)
1.3 环境安装
关于深度学习的环境安装可参考炮哥的利用Anaconda安装pytorch和paddle深度学习环境+pycharm安装—免额外安装CUDA和cudnn(适合小白的保姆级教学)，这里不再赘述。如果之前配置过 yolov5 的环境，yolov7 可直接使用。
2. 数据集和预训练权重的准备
2.1 数据集
这里训练采用的数据集是 PASCAL VOC 数据集，但博主并没有使用完整的 VOC 数据集，而是选用了部分数据，具体分布如下：
- 训练集：(VOC2007train + VOC2007val) x 80% = 4013
- 验证集：(VOC2007train + VOC2007val) x 20% = 998
- 测试集：0
这里给出下载链接 Baidu Drive【pwd:yolo】下载解压后整个数据集文件夹内容如下所示：
其中 images 存放训练集和验证集的图片文件，labels 存放着对应的 YOLO 格式的 .txt 文件。
完整的 VOC 数据集的相关介绍和下载可参考：目标检测：PASCAL VOC 数据集简介
由于大家可能从其它地方拿到的是 XML 格式的标签文件，这里提供一个 XML2YOLO 转换的代码，如下所示：(from ChatGPT)
import os
import cv2
import xml.etree.ElementTree as ET
import shutil
from multiprocessing import Pool, cpu_count
from tqdm import tqdm
import numpy as np
from functools import partial
def process_xml(xml_filename, img_path, xml_path, img_save_path, label_save_path, class_dict, ratio):
# 解析 xml 文件
xml_file_path = os.path.join(xml_path, xml_filename)
tree = ET.parse(xml_file_path)
root = tree.getroot()
# 获取图像的宽度和高度
img_filename = os.path.splitext(xml_filename)[0] + ".jpg"
img = cv2.imread(os.path.join(img_path, img_filename))
height, width = img.shape[:2]
# 随机决定当前图像和标签是属于训练集还是验证集
subset = "train" if np.random.random() < ratio else "val"
# 打开对应的标签文件进行写入
label_file = os.path.join(label_save_path, subset, os.path.splitext(xml_filename)[0] + ".txt")
with open(label_file, "w") as file:
for obj in root.iter('object'):
# 获取类别名并转换为类别ID
class_name = obj.find('name').text
class_id = class_dict[class_name]
# 获取并处理边界框的坐标
xmlbox = obj.find('bndbox')
x1 = float(xmlbox.find('xmin').text)
y1 = float(xmlbox.find('ymin').text)
x2 = float(xmlbox.find('xmax').text)
y2 = float(xmlbox.find('ymax').text)
# 计算中心点坐标和宽高，并归一化
x_center = (x1 + x2) / 2 / width
y_center = (y1 + y2) / 2 / height
w = (x2 - x1) / width
h = (y2 - y1) / height
# 写入文件
file.write(f"{class_id} {x_center} {y_center} {w} {h}\n")
# 将图像文件复制到对应的训练集或验证集目录
shutil.copy(os.path.join(img_path, img_filename), os.path.join(img_save_path, subset, img_filename))
def check_and_create_dir(path):
# 检查并创建 train 和 val 目录
for subset in ['train', 'val']:
if not os.path.exists(os.path.join(path, subset)):
os.makedirs(os.path.join(path, subset))
if __name__ == "__main__":
# 1. 定义路径和类别字典，不要使用中文路径
img_path = "D:\\Data\\PASCAL_VOC\\VOCdevkit\\VOC2007\\JPEGImages"
xml_path = "D:\\Data\\PASCAL_VOC\\VOCdevkit\\VOC2007\\Annotations"
img_save_path = "D:\\Data\\PASCAL_VOC\\dataset\\images"
label_save_path = "D:\\Data\\PASCAL_VOC\\dataset\\labels"
class_dict = {
"aeroplane": 0,
"bicycle": 1,
"bird": 2,
"boat": 3,
"bottle": 4,
"bus": 5,
"car": 6,
"cat": 7,
"chair": 8,
"cow": 9,
"diningtable": 10,
"dog": 11,
"horse": 12,
"motorbike": 13,
"person": 14,
"pottedplant": 15,
"sheep": 16,
"sofa": 17,
"train": 18,
"tvmonitor": 19
}
train_val_ratio = 0.8 # 2. 定义训练集和验证集的比例
# 检查并创建必要的目录
check_and_create_dir(img_save_path)
check_and_create_dir(label_save_path)
# 获取 xml 文件列表
xml_filenames = os.listdir(xml_path)
# 创建进程池并执行
with Pool(cpu_count()) as p:
list(tqdm(p.imap(partial(process_xml, img_path=img_path, xml_path=xml_path, img_save_path=img_save_path, label_save_path=label_save_path,
class_dict=class_dict, ratio=train_val_ratio), xml_filenames), total=len(xml_filenames)))
上述代码的功能是将 PASCAL VOC 格式的数据集（包括 JPEG 图像和 XML 格式的标签文件）转换为 YOLO 需要的 .txt 标签格式，同时会将转换后的数据集按照比例随机划分为训练集和验证集。
你需要修改以下几项：
- img_path：需要转换的图像文件路径
- xml_path：需要转换的 xml 标签文件路径
- img_save_path：转换后保存的图像路径
- label_save_path：转换后保存的 txt 标签路径
- class_dict：数据集类别字典
- train_val_ratio：训练集和验证集划分的比例
- 注意：以上路径都不要包含中文，Windows 下路径记得使用
\\或者/防止转义
XML 标签文件中目标框保存的格式是 [xmin, ymin, xmax, ymax] 四个变量，分别代表着未经归一化的左上角和右下角坐标。
YOLO 标签中目标框保存的格式是每一行代表一个目标框信息，每一行共包含 [label_id, x_center, y_center, w, h] 五个变量，分别代表着标签 ID，经过归一化后的中心点坐标和目标框宽高。
关于代码的分析可以参考：tensorRT模型性能测试
至此，数据集的准备工作完毕。
2.2 预训练权重准备
yolov7 预训练权重可以通过 here【pwd:yolo】下载，注意这是 yolov7-v0.1 版本的预训练权重，若后续有版本更新，记得替换。本次训练 VOC 数据集使用的预训练权重为 yolov7-tiny.pt
3. 训练模型
将准备好的数据集文件夹即 VOC 复制到 yolov7 项目环境中，将准备好的预训练权重 yolov7-tiny.pt 复制到 yolov7 项目环境中，完整的项目结构如下图所示。训练目标检测模型主要修改 cfg 文件夹下的模型配置文件 yolov7-tiny.ymal 以及 data 文件夹下的数据配置文件 coco.yaml
3.1 修改模型配置文件
由于该项目使用的是 yolov7-tiny.pt 这个预训练权重，所以需要使用 cfg/training 目录下的 yolov7-tiny.yaml 这个文件（由于不同的预训练权重对应不同的网络结构，所以用错预训练权重会报错）。主要修改 yolov7-tiny.yaml 文件的第二行，即需要识别的类别数，由于这里识别 VOC 的 20 个类别，故修改为 20 即可，如下所示
3.2 修改数据配置文件
修改 data 目录下相应的 yaml 文件，找到目录下的 coco.yaml 文件，主要修改如下：
- 1. 注释第 4 行
- 2. 修改第 7 行训练集的路径
- 3. 修改第 8 行验证集的路径
- 4. 注释第 9 行，因为未使用到测试集
- 5. 修改第 12 行需要检测的类别数个数
- 6. 修改第 15 行需要检测的类别数名称
3.3 训练模型
在终端执行如下指令即可开始训练，参考自 yolov7 的 README.md/Training
python train.py --workers 8 --device 0 --batch-size 32 --data data/coco.yaml --img 640 640 --cfg cfg/training/yolov7-tiny.yaml --weights 'yolov7-tiny.pt' --name yolov7 --hyp data/hyp.scratch.p5.yaml --epochs 100
博主训练的模型为 p5 models 且使用的是单个 GPU 进行训练，显卡为 RTX3060，操作系统为 Ubuntu20.04，pytorch 版本为 1.12.0，训练时长大概 1 小时左右。训练的参数的指定和 yolov5 差不多，简要解释如下：
- –-workers 最大工作核心数
- –-device 指定训练的设备，CPU，0(代表第一个 GPU 设备)
- –-batch-size 每次输入到网络的图片数
- -–data 数据配置文件的路径
- –-img 输入图像的尺寸
- –-cfg 模型配置文件路径
- –-weights 预训练权重路径
- –-name 训练保存的文件夹名字
- -–hyp 超参数文件路径
- –epochs 训练轮数
还有其它参数博主并未设置，如 –-multi-scale 多尺度训练等。大家一定要根据自己的实际情况(如显卡算力)指定不同的参数，如果你之前训练过 yolov5，那我相信这对你来说应该是小 case
训练完成后的模型权重保存在 run/train/weights 文件夹下，和 yolov5 不同的是它保存了多个权重文件，使用 best.pt 进行后续模型部署量化即可，这里提供博主训练好的权重文件下载链接 Baidu Drive【pwd:yolo】
3.4 mAP测试
由于后续我们要对模型进行 PTQ 量化，需要一些指标来衡量模型的性能，mAP 是一个重要的衡量指标。我们需要对比量化前后模型的 mAP，首先来看量化前原始 pytorch 模型的 mAP，测试的数据集直接选用验证集的 998 张图片。
我们将置信度阈值设置为 0.001，NMS 阈值设置为 0.65，方便与后续 PTQ 量化模型对比。
mAP 测试的指令如下：
python test.py --data data/coco.yaml --img 640 --batch 32 --conf 0.001 --iou 0.65 --device 0 --weights best.pt --name yolov7_640_val
测试完成后的结果会保存在 runs/test/yolov7_640_val 文件夹下，这里总结下原始 pytorch 模型的性能
| Model | Size | mAPval 0.5:0.95 |
mAPval 0.5 |
Params (M) |
FLOPs (G) |
|---|---|---|---|---|---|
| YOLOv7-tiny | 640 | 0.491 | 0.744 | 5.8 | 13.3 |
三、YOLOv7-QAT准备工作
参考自：https://github.com/NVIDIA-AI-IOT/yolo_deepstream/tree/main/yolov7_qat/README.md
描述：在正式开始 QAT 量化之前我们需要做一些准备工作，比如安装一些必要的依赖库，准备好用于量化训练的权重和数据集，以及简单修改部分代码等等。请大家务必熟读对应的 README 文档，将博主实现的流程对照着 README 文档来看可能更方便理解。
1. 项目克隆
克隆 yolov7 项目
git clone https://github.com/WongKinYiu/yolov7.git
克隆 yolo_deepstream 项目
git clone https://github.com/NVIDIA-AI-IOT/yolo_deepstream.git
也可以点击 here【pwd:yolo】 下载博主准备好的代码（注意该代码下载于 2023/10/21 日，若有改动请参考最新）
将 yolo_deepstream 项目中的 yolov7_qat 文件夹下的所有文件复制到 yolov7 项目中，指令如下：
cp -r yolo_deepstream/yolov7_qat/* yolov7/
完整的目录如下：
2. 安装依赖
QAT 量化需要使用到 NVIDIA 为 TensorRT 提供的 pytorch-quantization 工具，安装指令如下：
pip install pytorch-quantization --extra-index-url https://pypi.ngc.nvidia.com
3. 数据集和权重准备
我们需要准备一个数据集用于 QAT 模型的微调，数据集直接拿之前用于训练的数据集即可，此外我们还需要提供两个 txt 文档，一个是 train2017.txt 里面包含所有训练集图片的完整路径，一个是 val2017.txt 里面包含所有验证集图片的完整路径。
txt 文档生成代码如下：
import os
save_dir = "/home/jarvis/Learn/Datasets/VOC_QAT"
train_dir = "/home/jarvis/Learn/Datasets/VOC_QAT/images/train"
train_txt_path = os.path.join(save_dir, "train2017.txt")
with open(train_txt_path, "w") as f:
for filename in os.listdir(train_dir):
if filename.endswith(".jpg") or filename.endswith(".png"): # 添加你的图像文件扩展名
file_path = os.path.join(train_dir, filename)
f.write(file_path + "\n")
print(f"train2017.txt has been created at {train_txt_path}")
val_dir = "/home/jarvis/Learn/Datasets/VOC_QAT/images/val"
val_txt_path = os.path.join(save_dir, "val2017.txt")
with open(val_txt_path, "w") as f:
for filename in os.listdir(val_dir):
if filename.endswith(".jpg") or filename.endswith(".png"): # 添加你的图像文件扩展名
file_path = os.path.join(val_dir, filename)
f.write(file_path + "\n")
print(f"val2017.txt has been created at {val_txt_path}")
你需要修改以下几项：
- save_dir：txt 文档保存的路径，应该与 images 和 labels 文件夹在同一级目录
- train_dir：训练集图片路径
- val_dir：验证集图片路径
执行完成后会在对应目录下生成 train2017.txt 和 val2017.txt 两个文件。
数据集完整的目录结构如下：
.
├── images
│ ├── train
│ └── val
├── labels
│ ├── train
│ └── val
├── train2017.txt
└── val2017.txt
6 directories, 2 files
除数据集外我们还需要准备一个权重文件用于 QAT 量化训练，权重直接选取之前 yolov7 训练 VOC 数据集的 best.pt 文件即可。可以点击 here【pwd:yolo】下载博主准备好的数据集和权重。
我们可以将准备好的数据集和权重都放在 yolov7 项目下，方便后续操作。
4. 配置文件修改
由于 QAT 量化过程需要训练，因此我们还需要修改下 yolov7 目录下的配置文件方便后续训练。
主要修改数据配置文件 data/coco.yaml 以及模型配置文件 cfg/training/yolov7-tiny.yaml，我们在之前模型训练中有详细提到过，这边不再赘述。
至此，YOLOv7-QAT 的准备工作到这里就结束了，下面我们正式开始 QAT 量化训练和部署
四、YOLOv7-QAT微调导出
在正式开始 QAT 量化训练之前，请务必确保完成了三中的工作，这对我们后续的量化训练部署非常重要。
我们使用 TensorRT 的 pytorch quantization 量化工具对 yolov7-QAT 模型进行微调训练，然后将模型导出为 ONNX 并利用 TensorRT 部署。
1. QAT微调
将代码、数据集、权重准备好后我们就可以来进行 QAT 量化了，进入 yolov7 主目录执行如下指令：
python scripts/qat.py quantize best.pt --qat=qat.pt --ptq=ptq.pt --cocodir=/home/jarvis/Learn/Datasets/VOC_QAT --ignore-policy="model\.77\.m\.(.*)|model\.0\.(.*)" --supervision-stride=1 --eval-ptq --eval-origin
注：由于博主训练的是 yolov7-tiny 模型，因此它的 QAT 量化指令和 yolov7 是不同的，具体可参考 yolov7_qat/README.md
注意将 cocodir 替换成你自己的路径
该指令会利用 best.pt 权重和对应的数据集进行量化训练，量化训练过程如下图所示：
量化完成后在当前目录下会生成 ptq.pt 和 qat.pt 模型文件，分别对应着 PTQ 模型和 QAT 模型，后续我们只需要 qat.pt 模型并导出为 ONNX 即可。
上述 QAT 量化训练的步骤如下：
- 插入 Q/DQ 节点获得伪量化的 pytorch 模型。值得注意的是 pytorch quantization 工具提供了自动插入 Q/DQ 节点的功能。但对于 yolov7 模型，它无法获得与 PTQ 相同的性能，因为在显式模式（QAT 模式）下，TensorRT 会参考 Q/DQ 节点的放置的位置来限制模型的精度。一些自动添加的 Q/DQ 节点无法与其他层融合，这将导致一些额外的无用精度转换。在我们的脚本中，我们为 yolov7 找到了一些规则和限制，以基于规则的方式自动分析和配置 Q/DQ 节点，确保它们在 TensorRT 下达到最佳状态，并确保所有节点都以 INT8 精度运行，关于这部分的细节可以参考 quantization/rules.py 文件
- PTQ 校准。我们建议先运行 PTQ-Calibration，根据实验，Histogram(MSE) 是 yolov7 最佳 PTQ 校准方法。值得注意的是，如果你对 PTQ 结果满意，也可以跳过 QAT
- QAT 训练。QAT 需要微调训练我们的模型，当我们获得一个比较满意的精度时，可以将对应的权重保存下来
2. QAT模型导出
我们需要将上面生成的 qat.pt 导出为 ONNX，在导出之前我们需要修改下源代码让其导出的 ONNX 模型能够适配 tensorRT_Pro。为此我们需要修改 scripts/qat.py 文件，修改如下：
# scripts/qat.py第138行，export_onnx函数
# quantize.export_onnx(model, dummy, file, opset_version=13,
# input_names=["images"], output_names=["outputs"],
# dynamic_axes={"images": {0: "batch"}, "outputs": {0: "batch"}} if dynamic_batch else None
# )
# 修改为：
quantize.export_onnx(model, dummy, file, opset_version=13,
input_names=["images"], output_names=["output"],
dynamic_axes={"images": {0: "batch"}, "output": {0: "batch"}} if dynamic_batch else None
)
修改完成后我们就可以导出 qat.pt 模型了，指令如下：
python scripts/qat.py export qat.pt --size=640 --save=qat.onnx --dynamic
输出如下：
导出的 ONNX 模型如下图所示：
3. INT8模型生成
TensorRT8 的新特性是可以加载带有 QDQ 信息的模型然后生成对应量化版本的 engine。因此我们可以将上面导出的带有量化信息的 QAT 模型利用 trtexec 工具生成对应的 engine，对于 trtexec 有困惑的可以参考：如何熟练的使用trtexec
我们新建一个 build.sh 脚本文件来生成 engine，其内容如下：
#! /usr/bin/bash
echo "Build INT8 Model"
TRTEXEC=/opt/TensorRT-8.4.1.5/bin/trtexec
${TRTEXEC} --onnx=qat.onnx --minShapes=images:1x3x640x640 --optShapes=images:1x3x640x640 --maxShapes=images:16x3x640x640 --fp16 --int8 --saveEngine=qat.INT8.trtmodel
注意将 TRTEXEC 替换成你自己的路径
将 build.sh 脚本文件准备好后，我们可以在终端执行如下指令生成对应的 INT8 模型：
bash build.sh
输出如下图所示：
执行成功后会在当前目录下生成 qat.INT8.trtmodel 模型文件，拿到 engine 文件后就可以使用 tensorRT_Pro 完成后续的部署工作了。
至此，YOLOv7-QAT 的微调导出到这里就结束了。
可以点击 here【pwd:yolo】下载博主 QAT 量化训练好的 qat 模型和导出的 onnx 模型。
五、YOLOv7-QAT部署
由于博主手头没有合适的 Jetson 嵌入式设备，因此打算使用自己的主机完成 YOLOv7-QAT 部署工作，部署使用的 repo 是 tensorRT_Pro。
接下来我们主要是针对 tensorRT_Pro 项目中的 YOLOv7 完成 QAT 的模型部署，本次部署的模型是 YOLOv7-tiny.pt，数据集为 VOC，类别数为 20。
1. 源码下载
tensorRT_Pro 的代码可以直接从 GitHub 官网上下载，源码下载地址是 https://github.com/shouxieai/tensorRT_Pro，Linux 下代码克隆指令如下：
$ git clone https://github.com/shouxieai/tensorRT_Pro
也可手动点击下载，点击右上角的 Code 按键，将代码下载下来。至此整个项目就已经准备好了。也可以点击 Baidu Drive【pwd:yolo】 下载博主准备好的源代码（注意该代码下载于 2023/10/21 日，若有改动请参考最新）
2. 环境配置
需要使用的软件环境有 TensorRT、CUDA、cuDNN、OpenCV、Protobuf，所有软件环境的安装可以参考 Ubuntu20.04部署YOLOv5，这里不再赘述，需要各位看官自行配置好相关环境，外网访问较慢，这里提供下博主安装过程中的软件安装包下载链接 Baidu Drive【pwd:yolo】
tensorRT_Pro 提供 CMakeLists.txt 和 Makefile 两种方式编译，二者选一即可
2.1 配置CMakeLists.txt
主要修改六处
1. 修改第 10 行，选择不支持 python (也可选择支持)
set(HAS_PYTHON OFF)
2. 修改第 18 行，修改 OpenCV 路径
set(OpenCV_DIR "/usr/local/include/opencv4/")
3. 修改第 20 行，修改 CUDA 路径
set(CUDA_TOOLKIT_ROOT_DIR "/usr/local/cuda-11.6")
4. 修改第 21 行，修改 cuDNN 路径
set(CUDNN_DIR "/usr/local/cudnn8.4.0.27-cuda11.6")
5. 修改第 22 行，修改 tensorRT 路径
set(TENSORRT_DIR "/opt/TensorRT-8.4.1.5")
6. 修改第 33 行，修改 protobuf 路径
set(PROTOBUF_DIR "/home/jarvis/protobuf")
完整的 CMakeLists.txt 的内容如下：
cmake_minimum_required(VERSION 2.6)
project(pro)
option(CUDA_USE_STATIC_CUDA_RUNTIME OFF)
set(CMAKE_CXX_STANDARD 11)
set(CMAKE_BUILD_TYPE Debug)
set(EXECUTABLE_OUTPUT_PATH ${PROJECT_SOURCE_DIR}/workspace)
# 如果要支持python则设置python路径
set(HAS_PYTHON OFF) # ===== 修改 1 =====
set(PythonRoot "/datav/software/anaconda3")
set(PythonName "python3.9")
# 如果你是不同显卡，请设置为显卡对应的号码参考这里：https://developer.nvidia.com/zh-cn/cuda-gpus#compute
#set(CUDA_GEN_CODE "-gencode=arch=compute_75,code=sm_75")
# 如果你的opencv找不到，可以自己指定目录
set(OpenCV_DIR "/usr/local/include/opencv4/") # ===== 修改 2 =====
set(CUDA_TOOLKIT_ROOT_DIR "/usr/local/cuda-11.6") # ===== 修改 3 =====
set(CUDNN_DIR "/usr/local/cudnn8.4.0.27-cuda11.6") # ===== 修改 4 =====
set(TENSORRT_DIR "/opt/TensorRT-8.4.1.5") # ===== 修改 5 =====
# set(CUDA_TOOLKIT_ROOT_DIR "/data/sxai/lean/cuda-10.2")
# set(CUDNN_DIR "/data/sxai/lean/cudnn7.6.5.32-cuda10.2")
# set(TENSORRT_DIR "/data/sxai/lean/TensorRT-7.0.0.11")
# set(CUDA_TOOLKIT_ROOT_DIR "/data/sxai/lean/cuda-11.1")
# set(CUDNN_DIR "/data/sxai/lean/cudnn8.2.2.26")
# set(TENSORRT_DIR "/data/sxai/lean/TensorRT-7.2.1.6")
# 因为protobuf，需要用特定版本，所以这里指定路径
set(PROTOBUF_DIR "/home/jarvis/protobuf") # ===== 修改 6 ======
find_package(CUDA REQUIRED)
find_package(OpenCV)
include_directories(
${PROJECT_SOURCE_DIR}/src
${PROJECT_SOURCE_DIR}/src/application
${PROJECT_SOURCE_DIR}/src/tensorRT
${PROJECT_SOURCE_DIR}/src/tensorRT/common
${OpenCV_INCLUDE_DIRS}
${CUDA_TOOLKIT_ROOT_DIR}/include
${PROTOBUF_DIR}/include
${TENSORRT_DIR}/include
${CUDNN_DIR}/include
)
# 切记，protobuf的lib目录一定要比tensorRT目录前面，因为tensorRTlib下带有protobuf的so文件
# 这可能带来错误
link_directories(
${PROTOBUF_DIR}/lib
${TENSORRT_DIR}/lib
${CUDA_TOOLKIT_ROOT_DIR}/lib64
${CUDNN_DIR}/lib
)
if("${HAS_PYTHON}" STREQUAL "ON")
message("Usage Python ${PythonRoot}")
include_directories(${PythonRoot}/include/${PythonName})
link_directories(${PythonRoot}/lib)
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DHAS_PYTHON")
endif()
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++11 -Wall -O0 -Wfatal-errors -pthread -w -g")
set(CUDA_NVCC_FLAGS "${CUDA_NVCC_FLAGS} -std=c++11 -O0 -Xcompiler -fPIC -g -w ${CUDA_GEN_CODE}")
file(GLOB_RECURSE cpp_srcs ${PROJECT_SOURCE_DIR}/src/*.cpp)
file(GLOB_RECURSE cuda_srcs ${PROJECT_SOURCE_DIR}/src/*.cu)
cuda_add_library(plugin_list SHARED ${cuda_srcs})
target_link_libraries(plugin_list nvinfer nvinfer_plugin)
target_link_libraries(plugin_list cuda cublas cudart cudnn)
target_link_libraries(plugin_list protobuf pthread)
target_link_libraries(plugin_list ${OpenCV_LIBS})
add_executable(pro ${cpp_srcs})
# 如果提示插件找不到，请使用dlopen(xxx.so, NOW)的方式手动加载可以解决插件找不到问题
target_link_libraries(pro nvinfer nvinfer_plugin)
target_link_libraries(pro cuda cublas cudart cudnn)
target_link_libraries(pro protobuf pthread plugin_list)
target_link_libraries(pro ${OpenCV_LIBS})
if("${HAS_PYTHON}" STREQUAL "ON")
set(LIBRARY_OUTPUT_PATH ${PROJECT_SOURCE_DIR}/example-python/pytrt)
add_library(pytrtc SHARED ${cpp_srcs})
target_link_libraries(pytrtc nvinfer nvinfer_plugin)
target_link_libraries(pytrtc cuda cublas cudart cudnn)
target_link_libraries(pytrtc protobuf pthread plugin_list)
target_link_libraries(pytrtc ${OpenCV_LIBS})
target_link_libraries(pytrtc "${PythonName}")
target_link_libraries(pro "${PythonName}")
endif()
add_custom_target(
yolo
DEPENDS pro
WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}/workspace
COMMAND ./pro yolo
)
add_custom_target(
yolo_gpuptr
DEPENDS pro
WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}/workspace
COMMAND ./pro yolo_gpuptr
)
add_custom_target(
yolo_fast
DEPENDS pro
WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}/workspace
COMMAND ./pro yolo_fast
)
add_custom_target(
centernet
DEPENDS pro
WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}/workspace
COMMAND ./pro centernet
)
add_custom_target(
alphapose
DEPENDS pro
WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}/workspace
COMMAND ./pro alphapose
)
add_custom_target(
retinaface
DEPENDS pro
WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}/workspace
COMMAND ./pro retinaface
)
add_custom_target(
dbface
DEPENDS pro
WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}/workspace
COMMAND ./pro dbface
)
add_custom_target(
arcface
DEPENDS pro
WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}/workspace
COMMAND ./pro arcface
)
add_custom_target(
bert
DEPENDS pro
WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}/workspace
COMMAND ./pro bert
)
add_custom_target(
fall
DEPENDS pro
WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}/workspace
COMMAND ./pro fall_recognize
)
add_custom_target(
scrfd
DEPENDS pro
WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}/workspace
COMMAND ./pro scrfd
)
add_custom_target(
lesson
DEPENDS pro
WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}/workspace
COMMAND ./pro lesson
)
add_custom_target(
pyscrfd
DEPENDS pytrtc
WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}/example-python
COMMAND python test_scrfd.py
)
add_custom_target(
pyinstall
DEPENDS pytrtc
WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}/example-python
COMMAND python setup.py install
)
add_custom_target(
pytorch
DEPENDS pytrtc
WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}/example-python
COMMAND python test_torch.py
)
add_custom_target(
pyyolov5
DEPENDS pytrtc
WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}/example-python
COMMAND python test_yolov5.py
)
add_custom_target(
pycenternet
DEPENDS pytrtc
WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}/example-python
COMMAND python test_centernet.py
)
2.2 配置Makefile
主要修改六处
1. 修改第 4 行，修改 protobuf 路径
lean_protobuf := /home/jarvis/protobuf
2. 修改第 5 行，修改 tensorRT 路径
lean_tensor_rt := /opt/TensorRT-8.4.1.5
3. 修改第 6 行，修改 cuDNN 路径
lean_cudnn := /usr/local/cudnn8.4.0.27-cuda11.6
4. 修改第 7 行，修改 OpenCV 路径
lean_opencv := /usr/local
5. 修改第 8 行，修改 CUDA 路径
lean_cuda := /usr/local/cuda-11.6
6. 修改第 9 行，选择不支持 python (也可选择支持)
use_python := false
完整的 Makefile 的内容如下：
cc := g++
nvcc = ${lean_cuda}/bin/nvcc
lean_protobuf := /home/jarvis/protobuf # ===== 修改 1 =====
lean_tensor_rt := /opt/TensorRT-8.4.1.5 # ===== 修改 2 =====
lean_cudnn := /usr/local/cudnn8.4.0.27-cuda11.6 # ===== 修改 3 =====
lean_opencv := /usr/local # ===== 修改 4 =====
lean_cuda := /usr/local/cuda-11.6 # ===== 修改 5 =====
use_python := false # ===== 修改 6 =====
python_root := /datav/software/anaconda3
# python_root指向的lib目录下有个libpython3.9.so，因此这里写python3.9
# 对于有些版本，so名字是libpython3.7m.so，你需要填写python3.7m
# /datav/software/anaconda3/lib/libpython3.9.so
python_name := python3.9
# 如果是其他显卡，请修改-gencode=arch=compute_75,code=sm_75为对应显卡的能力
# 显卡对应的号码参考这里：https://developer.nvidia.com/zh-cn/cuda-gpus#compute
cuda_arch := # -gencode=arch=compute_75,code=sm_75
cpp_srcs := $(shell find src -name "*.cpp")
cpp_objs := $(cpp_srcs:.cpp=.cpp.o)
cpp_objs := $(cpp_objs:src/%=objs/%)
cpp_mk := $(cpp_objs:.cpp.o=.cpp.mk)
cu_srcs := $(shell find src -name "*.cu")
cu_objs := $(cu_srcs:.cu=.cu.o)
cu_objs := $(cu_objs:src/%=objs/%)
cu_mk := $(cu_objs:.cu.o=.cu.mk)
include_paths := src \
src/application \
src/tensorRT \
src/tensorRT/common \
$(lean_protobuf)/include \
$(lean_opencv)/include/opencv4 \
$(lean_tensor_rt)/include \
$(lean_cuda)/include \
$(lean_cudnn)/include
library_paths := $(lean_protobuf)/lib \
$(lean_opencv)/lib \
$(lean_tensor_rt)/lib \
$(lean_cuda)/lib64 \
$(lean_cudnn)/lib
link_librarys := opencv_core opencv_imgproc opencv_videoio opencv_imgcodecs \
nvinfer nvinfer_plugin \
cuda cublas cudart cudnn \
stdc++ protobuf dl
# HAS_PYTHON表示是否编译python支持
support_define :=
ifeq ($(use_python), true)
include_paths += $(python_root)/include/$(python_name)
library_paths += $(python_root)/lib
link_librarys += $(python_name)
support_define += -DHAS_PYTHON
endif
empty :=
export_path := $(subst $(empty) $(empty),:,$(library_paths))
run_paths := $(foreach item,$(library_paths),-Wl,-rpath=$(item))
include_paths := $(foreach item,$(include_paths),-I$(item))
library_paths := $(foreach item,$(library_paths),-L$(item))
link_librarys := $(foreach item,$(link_librarys),-l$(item))
cpp_compile_flags := -std=c++11 -g -w -O0 -fPIC -pthread -fopenmp $(support_define)
cu_compile_flags := -std=c++11 -g -w -O0 -Xcompiler "$(cpp_compile_flags)" $(cuda_arch) $(support_define)
link_flags := -pthread -fopenmp -Wl,-rpath='$$ORIGIN'
cpp_compile_flags += $(include_paths)
cu_compile_flags += $(include_paths)
link_flags += $(library_paths) $(link_librarys) $(run_paths)
ifneq ($(MAKECMDGOALS), clean)
-include $(cpp_mk) $(cu_mk)
endif
pro : workspace/pro
pytrtc : example-python/pytrt/libpytrtc.so
expath : library_path.txt
library_path.txt :
@echo LD_LIBRARY_PATH=$(export_path):"$$"LD_LIBRARY_PATH > $@
workspace/pro : $(cpp_objs) $(cu_objs)
@echo Link $@
@mkdir -p $(dir $@)
@$(cc) $^ -o $@ $(link_flags)
example-python/pytrt/libpytrtc.so : $(cpp_objs) $(cu_objs)
@echo Link $@
@mkdir -p $(dir $@)
@$(cc) -shared $^ -o $@ $(link_flags)
objs/%.cpp.o : src/%.cpp
@echo Compile CXX $<
@mkdir -p $(dir $@)
@$(cc) -c $< -o $@ $(cpp_compile_flags)
objs/%.cu.o : src/%.cu
@echo Compile CUDA $<
@mkdir -p $(dir $@)
@$(nvcc) -c $< -o $@ $(cu_compile_flags)
objs/%.cpp.mk : src/%.cpp
@echo Compile depends CXX $<
@mkdir -p $(dir $@)
@$(cc) -M $< -MF $@ -MT $(@:.cpp.mk=.cpp.o) $(cpp_compile_flags)
objs/%.cu.mk : src/%.cu
@echo Compile depends CUDA $<
@mkdir -p $(dir $@)
@$(nvcc) -M $< -MF $@ -MT $(@:.cu.mk=.cu.o) $(cu_compile_flags)
yolo : workspace/pro
@cd workspace && ./pro yolo
yolo_gpuptr : workspace/pro
@cd workspace && ./pro yolo_gpuptr
dyolo : workspace/pro
@cd workspace && ./pro dyolo
dunet : workspace/pro
@cd workspace && ./pro dunet
dmae : workspace/pro
@cd workspace && ./pro dmae
dclassifier : workspace/pro
@cd workspace && ./pro dclassifier
yolo_fast : workspace/pro
@cd workspace && ./pro yolo_fast
bert : workspace/pro
@cd workspace && ./pro bert
alphapose : workspace/pro
@cd workspace && ./pro alphapose
fall : workspace/pro
@cd workspace && ./pro fall_recognize
retinaface : workspace/pro
@cd workspace && ./pro retinaface
arcface : workspace/pro
@cd workspace && ./pro arcface
test_warpaffine : workspace/pro
@cd workspace && ./pro test_warpaffine
test_yolo_map : workspace/pro
@cd workspace && ./pro test_yolo_map
arcface_video : workspace/pro
@cd workspace && ./pro arcface_video
arcface_tracker : workspace/pro
@cd workspace && ./pro arcface_tracker
test_all : workspace/pro
@cd workspace && ./pro test_all
scrfd : workspace/pro
@cd workspace && ./pro scrfd
centernet : workspace/pro
@cd workspace && ./pro centernet
dbface : workspace/pro
@cd workspace && ./pro dbface
high_perf : workspace/pro
@cd workspace && ./pro high_perf
lesson : workspace/pro
@cd workspace && ./pro lesson
plugin : workspace/pro
@cd workspace && ./pro plugin
pytorch : pytrtc
@cd example-python && python test_torch.py
pyscrfd : pytrtc
@cd example-python && python test_scrfd.py
pyretinaface : pytrtc
@cd example-python && python test_retinaface.py
pycenternet : pytrtc
@cd example-python && python test_centernet.py
pyyolov5 : pytrtc
@cd example-python && python test_yolov5.py
pyyolov7 : pytrtc
@cd example-python && python test_yolov7.py
pyyolox : pytrtc
@cd example-python && python test_yolox.py
pyarcface : pytrtc
@cd example-python && python test_arcface.py
pyinstall : pytrtc
@cd example-python && python setup.py install
clean :
@rm -rf objs workspace/pro example-python/pytrt/libpytrtc.so example-python/build example-python/dist example-python/pytrt.egg-info example-python/pytrt/__pycache__
@rm -rf workspace/single_inference
@rm -rf workspace/scrfd_result workspace/retinaface_result
@rm -rf workspace/YoloV5_result workspace/YoloX_result
@rm -rf workspace/face/library_draw workspace/face/result
@rm -rf build
@rm -rf example-python/pytrt/libplugin_list.so
@rm -rf library_path.txt
.PHONY : clean yolo alphapose fall debug
# 导出符号，使得运行时能够链接上
export LD_LIBRARY_PATH:=$(export_path):$(LD_LIBRARY_PATH)
3. QAT模型推理
我们来利用 tensorRT 加载 QAT 模型生成的 engine 进行推理，看效果咋样
开始之前你需要将之前生成的 qat.INT8.trtmodel 引擎文件放到 tensorRT_Pro/workspace 文件夹下。
此外我们还需要简单修改下源码，yolo 模型推理代码在 src/application/app_yolo.cpp 中，我们就只需要修改这一个文件中的内容即可，源码修改较简单主要有以下几点：
- 1. app_yolo.cpp 177 行，TRT::Mode 修改为 INT8，“yolov7” 改成 “qat”
- 2. app_yolo.cpp 25 行，新增 voclabels 数组，添加 voc 数据集的类别名称
- 3. app_yolo.cpp 100 行，cocolabels 修改为 voclabels
具体修改如下：
test(Yolo::Type::V7, TRT::Mode::INT8, "qat") // 修改1 177行 "yolov7"改成"qat"
static const char *voclabels[] = {"aeroplane", "bicycle", "bird", "boat", "bottle",
"bus", "car", "cat", "chair", "cow",
"diningtable", "dog", "horse", "motorbike", "person",
"pottedplant", "sheep", "sofa", "train", "tvmonitor"}; // 修改2 25行新增代码，为自训练模型的类别名称
for(auto& obj : boxes){
...
auto name = voclabels[obj.class_label]; // 修改3 100行cocolabels修改为voclabels
...
}
修改完成后在终端执行如下指令即可：
make yolo
图解如下所示：
编译运行后会在 workspace 文件夹下生成 qat_Yolov7_INT8_result 文件夹，该文件夹下保存了推理的图片。
模型推理效果如下图所示：
4. QAT模型mAP测试
经过 QAT 量化训练后的模型性能到底怎么样呢？下面我们来测试下经过 QAT 量化后模型的 mAP
我们需要适当修改下对应 mAP 测试的代码，在 src/application/test_yolo_map.cpp 文件夹中，我们就只需要修改这一个文件中的内容即可，源码修改较简单主要有以下几点：
- 1. test_yolo_map.cpp 172 行，修改要测试的验证集文件夹路径
- 2. test_yolo_map.cpp 175 行，修改要测试的 INT8 模型，yolov5s 修改为 qat
- 3. test_yolo_map.cpp 176 行，TRT::Mode 修改为 INT8
- 4. test_yolo_map.cpp 125 行，将 save_to_json 函数简单修改下
修改后完整的 test_yolo_map.cpp 如下所示：
#include 上述代码会将 INT8 模型在验证集中所有图像的检测结果存储到一个 JSON 文件中，每个检测到的物体都被序列化为 JSON 格式信息，包括图像 ID、类别 ID、置信度和边界框坐标。后续我们就可以拿着这个预测结果的 JSON 文件和我们真实标签的 JSON 文件通过 COCO Python API 去计算 mAP 指标。
有以下几点需要注意：
- 博主将 JSON 文件中的 image_id 保存为一个字符串，考虑到图片命名的差异性
- 博主将 JSON 文件中的 category_id 直接保存为类别标签，没有做转换
- mAP 测试使用的 NMS_threshold = 0.65f，Conf_threshold = 0.001f 与 pytorch 保持一致
- 关于 mAP 的相关原理介绍可参考 目标检测mAP计算以及coco评价标准
将源码修改好后，直接在终端执行如下指令即可：
make test_yolo_map
图解如下所示：
运行成功后在 workspace 文件夹下会生成 qat.prediction.json 文件，该 JSON 文件中保存着 INT8 模型在验证集上的推理结果。
我们拿到了模型预测结果的 JSON 文件后，还需要拿到真实标签的 JSON 文件，但是现在我们只有验证集真实的 YOLO 标签文件，因此需要将 YOLO 标签转换为 JSON 文件，转换代码如下：(from chatGPT)
import os
import cv2
import json
import logging
import os.path as osp
from tqdm import tqdm
from functools import partial
from multiprocessing import Pool, cpu_count
def set_logging(name=None):
rank = int(os.getenv('RANK', -1))
logging.basicConfig(format="%(message)s", level=logging.INFO if (rank in (-1, 0)) else logging.WARNING)
return logging.getLogger(name)
LOGGER = set_logging(__name__)
def process_img(image_filename, data_path, label_path):
# Open the image file to get its size
image_path = os.path.join(data_path, image_filename)
img = cv2.imread(image_path)
height, width = img.shape[:2]
# Open the corresponding label file
label_file = os.path.join(label_path, os.path.splitext(image_filename)[0] + ".txt")
with open(label_file, "r") as file:
lines = file.readlines()
# Process the labels
labels = []
for line in lines:
category, x, y, w, h = map(float, line.strip().split())
labels.append((category, x, y, w, h))
return image_filename, {"shape": (height, width), "labels": labels}
def get_img_info(data_path, label_path):
LOGGER.info(f"Get img info")
image_filenames = os.listdir(data_path)
with Pool(cpu_count()) as p:
results = list(tqdm(p.imap(partial(process_img, data_path=data_path, label_path=label_path), image_filenames), total=len(image_filenames)))
img_info = {image_filename: info for image_filename, info in results}
return img_info
def generate_coco_format_labels(img_info, class_names, save_path):
# for evaluation with pycocotools
dataset = {"categories": [], "annotations": [], "images": []}
for i, class_name in enumerate(class_names):
dataset["categories"].append(
{"id": i, "name": class_name, "supercategory": ""}
)
ann_id = 0
LOGGER.info(f"Convert to COCO format")
for i, (img_path, info) in enumerate(tqdm(img_info.items())):
labels = info["labels"] if info["labels"] else []
img_id = osp.splitext(osp.basename(img_path))[0]
img_h, img_w = info["shape"]
dataset["images"].append(
{
"file_name": os.path.basename(img_path),
"id": img_id,
"width": img_w,
"height": img_h,
}
)
if labels:
for label in labels:
c, x, y, w, h = label[:5]
# convert x,y,w,h to x1,y1,x2,y2
x1 = (x - w / 2) * img_w
y1 = (y - h / 2) * img_h
x2 = (x + w / 2) * img_w
y2 = (y + h / 2) * img_h
# cls_id starts from 0
cls_id = int(c)
w = max(0, x2 - x1)
h = max(0, y2 - y1)
dataset["annotations"].append(
{
"area": h * w,
"bbox": [x1, y1, w, h],
"category_id": cls_id,
"id": ann_id,
"image_id": img_id,
"iscrowd": 0,
# mask
"segmentation": [],
}
)
ann_id += 1
with open(save_path, "w") as f:
json.dump(dataset, f)
LOGGER.info(
f"Convert to COCO format finished. Resutls saved in {save_path}"
)
if __name__ == "__main__":
# Define the paths
data_path = "/home/jarvis/Learn/Datasets/VOC_PTQ/images/val"
label_path = "/home/jarvis/Learn/Datasets/VOC_PTQ/labels/val"
class_names = ["aeroplane", "bicycle", "bird", "boat", "bottle", "bus",
"car", "cat", "chair", "cow", "diningtable", "dog", "horse",
"motorbike", "person", "pottedplant", "sheep", "sofa", "train", "tvmonitor"] # 类别名称请务必与 YOLO 格式的标签对应
save_path = "./val.json"
img_info = get_img_info(data_path, label_path)
generate_coco_format_labels(img_info, class_names, save_path)
上述代码的功能是将 YOLO 格式的数据集（包括图像文件和对应的 .txt 标签文件）转换成 COCO JSON 格式的标注。转换后的数据包括一个 JSON 标签文件，JSON 标签文件中包含了每个图像的所有物体的类别和边界框信息。
你需要修改以下几项：
- data_path：需要转换的图像文件路径
- label_path：需要转换的 txt 标签文件路径
- class_names：数据集的类别列表，请务必与 YOLO 标签的相对应
- save_path：转换后 JSON 文件保存的路径
- 注意：以上路径都不要包含中文，Windows 下路径记得使用
\\或者/防止转义
YOLO 标签中目标框保存的格式是每一行代表一个目标框信息，每一行共包含 [label_id, x_center, y_center, w, h] 五个变量，分别代表着标签 ID，经过归一化后的中心点坐标和目标框宽高。
JSON 文件中目标框保存的格式是 [x，y，w，h] 四个变量，分别代表着经过归一化的左上角坐标和目标框宽高。
关于代码的分析可以参考：tensorRT模型性能测试
至此，两个 JSON 文件都准备好了，一个是模型推理的预测结果，一个是真实结果。拿到两个 JSON 文件后我们就可以进行 mAP 测试了，具体代码如下：
from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval
# Run COCO mAP evaluation
# Reference: https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocoEvalDemo.ipynb
annotations_path = "val.json"
results_file = "qat.prediction.json"
cocoGt = COCO(annotation_file=annotations_path)
cocoDt = cocoGt.loadRes(results_file)
imgIds = sorted(cocoGt.getImgIds())
cocoEval = COCOeval(cocoGt, cocoDt, 'bbox')
cocoEval.params.imgIds = imgIds
cocoEval.evaluate()
cocoEval.accumulate()
cocoEval.summarize()
你需要修改以下几项：
- annotations_path：真实标签的 JSON 文件路径
- results_file：模型预测结果的 JSON 文件路径
执行后测试结果如下图所示：
我们将它与原始 pytorch 的模型放在一起进行对比下：
| Model | Size | mAPval 0.5:0.95 |
mAPval 0.5 |
Params (M) |
FLOPs (G) |
|---|---|---|---|---|---|
| YOLOv7-tiny | 640 | 0.491 | 0.744 | 5.8 | 13.3 |
| YOLOv7-tiny-INT8 | 640 | 0.487 | 0.721 | - | - |
可以看到相比于原始 pytorch 模型，QAT 量化后的模型 mAP 下降了近 2 个点
OK！至此 YOLOv7 模型的 QAT 量化到这里结束了，各位看官可以在自己的数据集测试下 QAT 量化后模型的性能。
六、讨论
1. 不同精度模型对比
QAT 量化的模型性能到底怎么样呢？与其它精度的模型相比有哪些优势又有哪些劣势呢？
这个小节我们就来看看不同精度的模型的性能对比，主要从 mAP 和速度两个方面衡量。博主测试了在同一个验证集上原始 pytorch 模型，FP32 模型，FP16 模型，INT8 模型的性能。
原始 pytorch 模型和 INT8 模型性能我们之前已经了解过了，下面我们来看看 FP32 模型和 FP16 模型的性能。
FP32模型：
FP16模型：
INT8模型：
值得注意的是，关于速度的测试我们之前似乎并没有提到，它具体是如何测试的呢？
其实在 inference_and_performance 函数中就有关于速度相关的测试，主要说明如下：
- 1. 输入分辨率 640x640
- 2. batch_size = 1
- 3. 图像预处理 + 推理 + 后处理
- 4. CUDA-11.6，cuDNN-8.4.0，TensorRT-8.4.1.5
- 5. NVIDIA RTX3060
- 6. 测试次数，100 次取平均，去掉 warmup
- 7. 测试代码：src/application/app_yolo.cpp
- 8. 测试图像 6 张，位于 workspace/inference
- 分辨率分别为：810x1080，500x806，1024x684，550x676，1280x720，800x533
- 9. 测试方式，加载 6 张图后，以原图重复 100 次不停的塞进去。让模型经历完整的图像的预处理，后处理
测试结果如下表所示：
| Model | Precision | mAPval 0.5:0.95 |
mAPval 0.5 |
Elapsed Time/ms | FPS |
|---|---|---|---|---|---|
| YOLOv7-tiny.pt | - | 0.491 | 0.744 | - | - |
| YOLOv7-tiny-FP32 | FP32 | 0.488 | 0.724 | 2.82 | 355.15 |
| YOLOv7-tiny-FP16 | FP16 | 0.489 | 0.725 | 1.26 | 792.23 |
| YOLOv7-tiny-INT8 | INT8 | 0.487 | 0.721 | 1.48 | 675.67 |
可视化图如下所示：
从表中的数据我们可以分析得到下面的一些结论：
1. 精度与模型性能的关系
- 当我们从原始 pytorch 模型转到 FP32 模型时，正常来说应该基本是无损的，但是 mAP 掉了将近 2 个点左右，这并不符合我们的直觉。
- mAP 差 2 个点的原因可能是在固定分辨率这件事上，tensorRT 将图片分辨率固定在 640x640 大小。还有就是 pytorch 实现的所有细节并未完全加入进来，这些细节可能有没有找到的部分。
- FP32 模型和 FP16 模型的 mAP 几乎一样，没有任何精度的损失，这倒是符合我们的直觉
- QAT 量化训练后的 INT8 模型的 mAP 和 FP32 模型接近，这点倒是比 YOLOv5-QAT 量化训练后的效果要好
2. 速度与模型性能的关系
- FP16 和 INT8 的 FPS 分别为 792.23 和 675.67，远高于 FP32 的 355.15
- INT8 模型的速度还没有 FP16 快，那这可能是 tensorRT 内部在融合 Q/DQ 层时带来了一些额外的操作，或者也有可能硬件在处理经过层融合的模型结构时不那么高效
3. 权衡速度与精度
- FP32 提供了较好的精度，但速度较慢
- FP16 提供了与 FP32 类似的精度，但速度提高了约 2.3 倍，是一个非常不错的选择。
- QAT 量化训练后的 INT8 模型虽然保住了精度，但速度没提升呀，和 FP16 模型相比也没有任何优势
综上所述，在实际应用中，需要根据具体的需求权衡速度和精度。例如，对于实时应用，可能会选择 FP16 或 INT8 以获得更高的速度，尽管可能牺牲一些精度。而对于需要高精度的应用，可能会选择 FP32。
2. PTQ vs. QAT
PTQ 量化和 QAT 量化生成的 INT8 模型哪个更好呢？我们该选择哪种量化方式呢？
这个小节我们就来看看 PTQ 量化和 QAT 量化后的 INT8 模型性能对比，主要从 mAP 和速度两个方面衡量。
值得注意的是，博主在上篇文章使用 PTQ 量化的 pytorch 模型和本篇文章使用 QAT 量化的 pytorch 模型是同一个，因此我们直接把上篇文章的结果拿过来对比就行。
测试结果如下表所示：
| Model | Method | mAPval 0.5:0.95 |
mAPval 0.5 |
Elapsed Time/ms | FPS |
|---|---|---|---|---|---|
| YOLOv7-tiny | PTQ | 0.473 | 0.705 | 0.94 | 1066.55 |
| YOLOv7-tiny | QAT | 0.487 | 0.721 | 1.48 | 675.67 |
可视化图如下：
从表中的数据我们可以分析得到下面的结论：
1. 量化方式与精度的关系
- QAT 与 PTQ 量化方式生成的 INT8 模型在 mAP 上的差异相对较大，QAT 模型要好
2. 量化方式与速度的关系
- PTQ 量化生成的 INT8 模型明显比 QAT 量化生成的 INT8 模型推理速度更快
3. 量化方式的选择
- PTQ 量化作为一个后训练量化方法，提供比 QAT 更快的速度，其实现相对简单，因为它不需要训练过程，只需要准备校准图片即可。
- QAT 量化尽管其实现可能比 PTQ 更复杂，速度也较慢，但提供了好一点的 mAP。因此，如果追求更高的精度，并且愿意投入更多的时间和资源进行训练，QAT 量化是一个更好的选择。
QAT 实际效果和博主想象的还是差不多的，QAT 量化训练后的 INT8 模型的 mAP 与 FP32 模型相当，但这速度属实有点接受不了，也不清楚是不是博主某些操作没做或者没做对。
但是有一个比较好的点就是 QAT 量化训练后生成的 INT8 模型可以手动去控制每一层的精度，不至于完全不可控。另外也不需要去考虑 PTQ 模型的校准图片数量选取，某种意义上是进阶版的 PTQ，只是实现的方式和流程略微复杂了点。
3. YOLOv5-QAT vs. YOLOv7-QAT
最后我们当然是来对比下 YOLOv5-QAT 量化后模型的性能和 YOLOv7-QAT 量化后模型的性能哪个会更好，那其实两个模型训练用的数据集都是同一个啦，所以还是有可比性的
结果对比如下表所示：
| Model | Method | mAPval 0.5:0.95 |
mAPval 0.5 |
Elapsed Time/ms | FPS |
|---|---|---|---|---|---|
| YOLOv5s.pt | QAT(一) | 0.412 | 0.660 | 0.99 | 1006.86 |
| YOLOv5s.pt | QAT(二) | 0.423 | 0.667 | 0.99 | 1007.13 |
| YQOLOv7-tiny | QAT | 0.487 | 0.721 | 1.48 | 675.67 |
可视化图如下：
从表中我们可以看到对博主当前的 VOC 数据集而言，YOLOv7-tiny-QAT 模型精度上比 YOLOv5s-QAT 模型要好上不少，但速度上却差了很多。值得一提的是，这里对比的 YOLOv5s-QAT 模型都是在 tensorRT_Pro 中通过 C++ API 生成的，如果通过 trtexec 生成其速度会和 YOLOv7-tiny 差不多。
我们再来回顾下之前的 YOLOv5-QAT 量化的方案，首先是进行 QAT 量化训练，然后将带有 Q/DQ 节点的模型导出为 ONNX，随后我们对带有 Q/DQ 节点的 ONNX 模型进行了转换，将其转换成了对应的 PTQ 模型和校准缓存文件，最后利用校准缓存文件和对应的 PTQ ONNX 模型来生成 INT8 engine。
而对比这里的 YOLOv7-QAT 量化的方案，我们在 YOLOv7-QAT 中并没有做节点的转换，将其转换成 PTQ 模型，而是将导出的带有 Q/DQ 节点的 ONNX 模型直接通过 tensorRT 加载生成对应的 INT8 engine，那么在 YOLOv5-QAT 量化方案中的 Q/DQ 节点转换这个步骤有没有必要呢？它的目的是什么呢？
另外关于 QAT 量化训练后 INT8 模型的生成目前有以下几种方式：
- 1. 经过 QAT 量化训练后导出的带有 Q/DQ 的 ONNX 模型，通过 trtexec 直接生成 INT8，其 FPS 和 FP16 模型接近，YOLOv7-QAT 方案
- 2. 经过 QAT 量化训练后导出的带有 Q/DQ 的 ONNX 模型通过节点转换生成对应的 PTQ 模型和校准文件，校准文件和 PTQ ONNX 模型通过 trtexec 生成 INT8，其 FPS 和 FP16 模型接近，YOLOv5-QAT 方案
- 3. 经过 QAT 量化训练后导出的带有 Q/DQ 的 ONNX 模型通过节点转换生成对应的 PTQ 模型和校准文件，校准文件和 PTQ ONNX 模型通过 tensorRT_Pro 中的 C++ API 生成 INT8，其 FPS 比 FP16 快不少，YOLOv5-QAT 方案
具体造成这几种 INT8 模型速度差异的原因博主也并未发现，那最终我们想要的 INT8 模型肯定是速度和精度都要考虑的，但经过 YOLOv5-QAT 和 YOLOv7-QAT 的量化方案的测试结果下来，博主并未得到一个较好的结果，当然，博主在实验过程中肯定有些步骤没做或者没做对可能导致最终结果的偏差，各位看官若发现问题，请批评指正
OK！YOLOv7-QAT 量化的内容到这里就结束了，各位看官可以自行测试。
结语
本篇博客介绍了关于 yolov7 的 QAT 量化以及部署流程，博主在这里只做了最基础的演示，有些实现可能并没有完全按照 yolo_deepstream 来做，各位看官感兴趣的话可以自行测试。感谢各位看到最后，创作不易，读后有收获的看官帮忙点个⭐️
下载链接
- 软件安装包下载链接【提取码:yolo】
- 源代码、权重、数据集下载链接【提取码:yolo】
参考
- COCO Python API
- YOLOv5-QAT量化部署
- YOLOv7-PTQ量化部署
- tensorRT模型性能测试
- 如何熟练的使用trtexec
- Ubuntu20.04部署YOLOv5
- TensorRT量化第四课：PTQ与QAT
- 目标检测mAP计算以及coco评价标准
- 目标检测：PASCAL VOC 数据集简介
- https://github.com/WongKinYiu/yolov7
- https://github.com/shouxieai/tensorRT_Pro
- https://github.com/NVIDIA-AI-IOT/yolo_deepstream
- 利用Anaconda安装pytorch和paddle深度学习环境+pycharm安装—免额外安装CUDA和cudnn(适合小白的保姆级教学)