cap5：YoloV5分割任务的TensorRT部署指南（python版）
《TensorRT全流程部署指南》专栏文章目录：
- cap1：TensorRT介绍及CUDA环境安装
- cap2：1000分类的ResNet的TensorRT部署指南（python版）
- cap3：自定义数据集训练ResNet的TensorRT部署指南（python版）
- cap4：YoloV5目标检测任务的TensorRT部署指南（python版）
- cap5：YoloV5分割任务的TensorRT部署指南（python版）
…
文章目录
- 1、获取pth模型数据
- 2、导出ONNX
- 2.1 直接导出
- 2.2 测试
- 3、TensoRT环境搭建
- 4、转换TensorRT引擎
- 4.1 使用trtexec工具完成序列化
- 4.2 使用python的API进行转换
- 5、推理
- 5.1 完整代码
在前几章中，我们详细探讨了如何将分类模型ResNet和目标检测模型YOLOv5部署到TensorRT上，相信大家对TensorRT的模型部署已经有了初步的了解。接下来，本文将带领你进一步深入，探索如何将YOLOv5分割模型成功部署到TensorRT。
分割模型的部署流程与之前的分类和目标检测模型相似，唯一不同的是后处理会相对复杂一些。不过，值得注意的是，TensorRT的推理过程仅涉及模型推理，模型输出后的后处理任务由Python完成，和TensorRT无关。对于有一定经验的你来说，这些步骤将会变得轻松易掌握。接下来，我们将逐步解析这些关键环节，帮助你顺利掌握YOLOv5分割任务在TensorRT上的部署技巧。
1、获取pth模型数据
yolov5项目的官方地址在ultralytics/yolov5，这个项目提供了完整的yolo训练框架。使用这个框架可以训练出自己的检测模型。为了方便展示，本文直接使用官方提供的yolov5s-seg.pt预训练模型，地址在：Segmentation。官方的预训练模型的输入是(1,3,640,640)。
2、导出ONNX
2.1 直接导出
实际上导出onnx的脚本在ultralytics/yolov5/export.py中提供了，通过命令可以快速导出onnx模型：
python export.py --weights=yolov5s-seg.pt --imgsz=[640,640] --include==onnx
通过以上命令就可以导出onnx模型。export.py还可以导出更多类型，并定制一些需求，可以浏览export.py学习更多。
2.2 测试
首先，我们通过netron可视化onnx模型。可以看到该分割模型的输入为[1,3,640,640]，有两个输出：output0和output1
为了方便后续进行后处理（虽然和TensorRT无关），需要先弄清楚输出数据的含义。可以看到共计有25200个检测框，每个检测框有117的属性：4个bbox数据，1个bbox置信度，80个类别得分，32个mask分割掩码，共计117个值。
output1其实是原型掩码（protos），产生用于分割模型的原型掩码。
通俗来讲，原型掩码相当于最基础的32组合，32个mask就是决定使用哪些原型掩码并组合起来就得到当前bbox的分割结果。
3、TensoRT环境搭建
参考《cap2：1000分类的ResNet的TensorRT部署指南（python版）》中的3、环境搭建部分。
4、转换TensorRT引擎
这一步主要完成红框内的工作，就是将onnx模型序列化为TensorRT的engine模型。在这个过程中会自动根据设备的特异性进行优化，该过程十分缓慢。但是完成后会保存为本地文件，下次可以直接加载使用，不用再次序列化。
这个过程可以使用基于python的API完成，或者直接使用trtexec工具。
4.1 使用trtexec工具完成序列化
在第3节中下载了TensorRT包，在bin中有trtexec工具。打开终端进入trtexec工具所在文件夹就可以使用该命令工具了。
转换命令为：
trtexec --onnx=yolov5s-seg.onnx --saveEngine=yolov5s-seg.engine --fp16
onnx指定onnx模型路径，saveEngine知道转换后engine模保存路径，fp16表示使用FP16，这个可以大大提高推理速度。不加–fp16则默认使用FLOAT32
最后会出现PASSED结果，说明成功了，此时就得到engine模型。
4.2 使用python的API进行转换
使用python的API进行转换具有比较固定的流程，只不过可以根据需求比如动态输入等进行相应的设置。最终结果和trtexec转换没有什么区别。
不管什么方式转换得到的engine都可以被python或者c++拿去部署使用，但是要注意：1）版本对应：即转换使用的TensorRT版本和部署推理使用的版本需要一致；2）不支持跨设备：即在哪种设备上转换的就只能在哪种设备上使用，因为转换时根据设备特异性进行了优化。
import tensorrt as trt
# import pycuda.driver as cuda
# import pycuda.autoinit
# TensorRT 需要一个 日志对象，用于输出 WARNING 级别的日志，帮助调试问题
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
def build_engine(onnx_file_path, engine_file_path):
with trt.Builder(TRT_LOGGER) as builder, \
builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) as network, \
trt.OnnxParser(network, TRT_LOGGER) as parser:
# config 配置 TensorRT 编译选项
builder_config = builder.create_builder_config()
builder_config.set_flag(trt.BuilderFlag.FP16) # 设置FP16加速（如果 GPU 支持），提高计算速度并减少显存占用。注释则默认使用FP32
builder_config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30) # 设置最大工作空间（1GB），用于存储中间 Tensor 和计算资源
# 读取 ONNX 文件
with open(onnx_file_path, 'rb') as model:
if not parser.parse(model.read()):
for error in range(parser.num_errors):
print(parser.get_error(error))
return None
# 构建 TensorRT 引擎
serialized_engine = builder.build_serialized_network(network, builder_config)
# 开始反序列化为可执行引擎
runtime = trt.Runtime(TRT_LOGGER)
engine = runtime.deserialize_cuda_engine(serialized_engine)
# 序列化并保存引擎
with open(engine_file_path, 'wb') as f:
f.write(engine.serialize())
print(f"Saved TensorRT engine to {engine_file_path}")
if __name__ == "__main__":
onnx_file = "yolov5s-seg.onnx"
engine_file = "yolov5s-seg.engine"
build_engine(onnx_file, engine_file)
通过这种方式，我们也成功得到的engine模型。
5、推理
5.1 完整代码
推理代码也十分简单，先给出全部代码：
# 即将更新
推理代码其实非常简单，就是TensorRT部署的代码+python实现的预处理+python实现的后处理。只不过后处理比较麻烦，占了绝大多数的行数而已。其实只要掌握TensorRT推理的核心代码即可。