预计阅读时间：5 分钟
yolov5 fp16 动态输入转 TensorRT 部署到 triton 服务
把 yolov5 训练出的 PyTorch 模型转换成 TensorRT 模型，扔进 triton 里部署。一条看似官方全程护航的流水线，实际每一步都有暗雷，本文按流程逐个排掉。
需求
- 为加快推理速度，把 fp32 精度的模型转成 fp16；
- 为支持 yolov5 动态 padding 和批推理，输入的 batch 和图片 size 需要是动态的。
版本
- yolov5：v6.1
- triton docker 镜像：22.08
- TensorRT：8.4.2.4
- cudnn：8.2.1.32
- CUDA Toolkit：11.3.1
流程
1. pt 转 onnx
首先把 PyTorch 模型转成 onnx。yolov5 虽然自带转 engine 的方法，但 v6.1 版本里 --dynamic（动态输入）和 --half（fp16）两个参数不兼容，一起用会报错 --half not compatible with --dynamic.：
因此需要先单独转出 onnx，转换命令：
python export.py --weights last.pt --img 640 --dynamic --batch 1 --opset 13 --include onnx
注意 TensorRT 8 以上 onnx 算子库（opset）至少要指定为 13——默认是 12，用默认值转 engine 时会出问题。
转完去 onnx 可视化看一眼产物：
输入输出的 batch 和 size 都是动态的、不再是固定值，这一步没问题。
2. fp32 onnx 转 fp16 onnx
用 onnxmltools 转换：
import onnxmltools
from onnxmltools.utils.float16_converter import convert_float_to_float16
# Update the input name and path for your ONNX model
input_onnx_model = 'last.onnx'
# Change this path to the output name and path for your float16 ONNX model
output_onnx_model = 'last.onnx'
# Load your model
onnx_model = onnxmltools.utils.load_model(input_onnx_model)
# Convert tensor float type from your input ONNX model to tensor float16
onnx_model = convert_float_to_float16(onnx_model)
# Save as protobuf
onnxmltools.utils.save_model(onnx_model, output_onnx_model)
import onnx
model_onnx = onnx.load('last.onnx')
#onnx.checker.check_model(model_onnx)
print(model_onnx.graph.output)
再可视化一次，新模型的输入输出变为 float16，即转换成功。
3. onnx 转 TensorRT engine 模型
和第 1 步同样的问题：--dynamic 和 --half 不能同时用。对策是命令行只指定 --dynamic，--half 的部分进代码里改：
在 yolov5 export.py 中找到 export_engine 函数，把原 pt 模型和转出来的 fp16 onnx 模型放到同一路径下（文件名需要相同），注释掉 export_engine 里转 onnx 的部分，增加两行代码指定半精度转换：
再在 export_engine 中找到 dynamic 判断的部分，把 profile.set_shape 的维度改成自己的值。这个方法的四个参数分别是输入名、最小维度、最佳维度、最大维度：输入允许在最小与最大之间动态变化，最佳维度随便取个中间值即可，一般直接和最大维度写一样。
如果改用 TensorRT 自带的 trtexec 工具转，对应的参数是 --minShapes --optShapes --maxShapes，详见 trtexec 官方文档。
转换命令：
python export.py --weights last.pt --img 640 --dynamic --batch 1 --include engine --device 0 --verbose
成功会显示 success。如果 TensorRT 版本是 8.2 会报一个错误：
原因是 Range 算子要 TensorRT 8.4 才支持，升级 TensorRT 版本即可解决。
4. engine 模型部署到 triton
把转出来的模型重命名为 model.plan，放到 模型/1/model.plan。
镜像要选 22.08——转换用的 TensorRT 版本必须和 triton docker 镜像内置的 TensorRT 版本一致，各镜像内置的软件版本见 triton 镜像版本对应关系。
用下面的命令查看模型配置（your_model_name 换成自己模型的名字，8000 是 triton 的 http 端口）：
curl localhost:8000/v2/models/your_model_name/config
关于 batch 维度：如果在 config.pbtxt 里指定了 max_batch_size，batch 在 triton 中是隐式的，不会显式出现在输入输出维度里，input 和 output 的动态维度用 -1 表示、不带 batch 维；如果不加 max_batch_size，则输入输出最前面需要自己用 -1 指定动态 batch 维度。
5. 其他问题
模型转换的显卡系列和部署的显卡系列必须相同。在 30 系显卡上转换、拿到 20 系显卡上部署，会报算力不匹配的错误——20 系是 compute 7.5，30 系是 compute 8.6。
本文由 aboom 原创，转载请注明出处。