Nvidia在Jetson NX上模型量化

Nvidia 模型量化

整体流程：

需求：在Jetson NX 边缘计算盒上进行模型量化，仅处理yolov5检测模型。

使用前准备：1. Jetson计算盒。2. x86的linux服务器。3 .pt模型。

使用项目：https://github.com/wang-xinyu/tensorrtx/tree/yolov5-v7.0/yolov5

需要两个环境的原因是，生成wts需要torch环境，而Jetson是arm环境，不方便安装torch。

项目准备：

1. 确定训练.pt模型的yolov5版本，在yolov5的README文件中可以查看。当前支持yolov5 v1.0、v2.0、v3.0、v3.1、v4.0、v5.0、v6.0、v6.2、v7.0。
2. 下载.pt对应的yolov5官方代码：git clone -b v7.0 https://github.com/ultralytics/yolov5.git。
3. 下载与之对应的tensorrtx项目代码：git clone -b yolov5-v7.0 https://github.com/wang-xinyu/tensorrtx.git。

项目执行：

1. 生成.wts中间件(建议在x86_64的服务器中操作)。

   cd yolov5/   # ultralytics官方yolov5项目
   # wget https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5s.pt
   cp [tensorrtx]/yolov5/gen_wts.py .  # 拷贝tensorrtx项目中的脚本至yolov5下
   python gen_wts.py -w yolov5s.pt -o yolov5s.wts   # 生成.wts中间件
   # A file 'yolov5s.wts' will be generated.
   

2. 构建生成yolov5s.engine(需要在Jetson盒子中生成)。

   cd [tensorrtx]/yolov5/   
   # Update kNumClass in src/config.h if your model is trained on custom dataset
   mkdir build
   cd build
   # cp [ultralytics-yolov5]/yolov5s.wts . # 👿将刚才生成的wts文件拷贝至tensorrtx的build中，由于不是一个环境，此处需要手动
   cmake ..
   make
   
   # 👿将src/config.h中将constexpr static int kNumClass = 80;更改为模型的实际类别数，否则转化为engine会报错,更改后再次make
   make
   
   # 生成.engine：./yolov5_det -s yolov5s.wts yolov5s.engine s
   ./yolov5_det -s [.wts] [.engine] [n/s/m/l/x/n6/s6/m6/l6/x6 or c/c6 gd gw]  // serialize model to plan file
   # 验证.engine效果：./yolov5_det -d yolov5s.engine ../images 
   ./yolov5_det -d [.engine] [image folder]  // deserialize and run inference, the images in [image folder] will be processed.
   

3. INT8量化(前面使用FP16)

   • 准备构建量化表的图像。此处的图像需要来自于模型的训练集，200张左右即可，以文件夹的形式放至第二步中的build路径下即可。

   • src/config.h 中将宏由#define USE_FP16更改为USE_INT8并重新 make.

   • make
     
     # 重新生成INT8的.engine：./yolov5_det -s yolov5s.wts yolov5s.engine s
     ./yolov5_det -s [.wts] [.engine] [n/s/m/l/x/n6/s6/m6/l6/x6 or c/c6 gd gw]  // serialize model to plan file
     # 验证INT8.engine效果：./yolov5_det -d yolov5s.engine ../images 
     ./yolov5_det -d [.engine] [image folder]  // deserialize and run inference, the images in [image folder] will be processed.
     

说明

 1. 需要下载官方yolov5项目以及tensorrtx项目，其二者要对应，均需与.pt的版本对应。
 2. 经实际测试，量化后的效果较差。
 3. 仅尝试处理检测模型。