cap5：YoloV5分割任务的TensorRT部署指南（python版）

《TensorRT全流程部署指南》专栏文章目录：
  

• 《TensorRT全流程部署指南》专栏主页
  
• cap1：TensorRT先容及CUDA环境安装
  
• cap2：1000分类的ResNet的TensorRT部署指南（python版）
  
• cap3：自定义数据集训练ResNet的TensorRT部署指南（python版）
  
• cap4：YoloV5目标检测任务的TensorRT部署指南（python版）
  
• cap5：YoloV5分割任务的TensorRT部署指南（python版）
  …
  

  
  
在前几章中，我们详细探究了如何将分类模子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模子：

1. 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无关），必要先弄清楚输出数据的含义。可以看到输出实际上有两个：output0和output1。认识分割的同学自然知道如何从这两个输出获取bbox和分割mask。
output0输出共计有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模子。如下图所示，还会给出一些推理时间、内存占用等测试信息。

更多trtexec使用方法，来自DeepSeek的先容：
   trtexec 是 NVIDIA TensorRT 的命令行工具，用于优化和运行深度学习模子。它支持模子转换、性能测试和推理执行，常用于生产环境中的模子部署。
  常用功能

  

• 模子转换：将 ONNX、Caffe 等格式的模子转换为 TensorRT 引擎。
  
• 性能测试：评估模子在不同设置下的推理性能。
  
• 推理执行：使用 TensorRT 引擎进行推理。
  

  常用参数

  

• --onnx=<path>：指定 ONNX 模子路径。
  
• --deploy=<path>：指定 Caffe 模子路径。
  
• --output=<name>：指定输出节点名称。
  
• --batch=<N>：设置批量大小。
  
• --workspace=<size>：设置 GPU 工作空间大小（MB）。
  
• --fp16：启用 FP16 精度。
  
• --int8：启用 INT8 精度。
  
• --verbose：启用详细日志输出。
  
• --saveEngine=<path>：生存生成的 TensorRT 引擎。
  
• --loadEngine=<path>：加载已有的 TensorRT 引擎。
  
• --useDLACore=<N>：指定 DLA 核心（实用于 Jetson 装备）。
  

  示例

  

• 转换 ONNX 模子：
  
  1. trtexec --onnx=model.onnx --saveEngine=model.engine

  复制代码
• 性能测试：
  
  1. trtexec --loadEngine=model.engine --batch=16 --fp16

  复制代码
• 使用 INT8 精度：
  
  1. trtexec --onnx=model.onnx --int8 --batch=32

  复制代码

  总结

  trtexec 是 TensorRT 的重要工具，支持模子转换、性能测试和推理执行，通过调整参数可以优化模子性能。
  4.2 使用python的API进行转换

除了使用trtexec工具进行转换，也可以使用python的API进行转换，但是二者没有什么区别。
使用python的API进行转换具有比力固定的流程，只不过可以根据需求比如动态输入等进行相应的设置。
不管什么方式转换得到的engine都可以被python或者c++拿去部署使用，但是要注意：1）版本对应：即转换使用的TensorRT版本和部署推理使用的版本必要同等；2）不支持跨装备：即在哪种装备上转换的就只能在哪种装备上使用，因为转换时根据装备特异性进行了优化。

1. import tensorrt as trt
2. # import pycuda.driver as cuda
3. # import pycuda.autoinit
5. # TensorRT 需要一个 日志对象，用于输出 WARNING 级别的日志，帮助调试问题
6. TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
8. def build_engine(onnx_file_path, engine_file_path):
9.     with trt.Builder(TRT_LOGGER) as builder, \
10.          builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) as network, \
11.          trt.OnnxParser(network, TRT_LOGGER) as parser:
13.         # config 配置 TensorRT 编译选项
14.         builder_config = builder.create_builder_config()
15.         builder_config.set_flag(trt.BuilderFlag.FP16)   # 设置FP16加速（如果 GPU 支持），提高计算速度并减少显存占用。注释则默认使用FP32
16.         builder_config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)  # 设置最大工作空间（1GB），用于存储中间 Tensor 和计算资源
18.         # 读取 ONNX 文件
19.         with open(onnx_file_path, 'rb') as model:
20.             if not parser.parse(model.read()):
21.                 for error in range(parser.num_errors):
22.                     print(parser.get_error(error))
23.                 return None
25.         # 构建 TensorRT 引擎
26.         serialized_engine = builder.build_serialized_network(network, builder_config)
28.         # 开始反序列化为可执行引擎
29.         runtime = trt.Runtime(TRT_LOGGER)
30.         engine = runtime.deserialize_cuda_engine(serialized_engine)
32.         # 序列化并保存引擎
33.         with open(engine_file_path, 'wb') as f:
34.             f.write(engine.serialize())
35.         print(f"Saved TensorRT engine to {engine_file_path}")
37. if __name__ == "__main__":
38.     onnx_file = "yolov5s-seg.onnx"
39.     engine_file = "yolov5s-seg.engine"
40.     build_engine(onnx_file, engine_file)

复制代码

通过这种方式，我们也成功得到的engine模子。
5、推理

5.1 完整代码

推理代码也十分简单，重要流程就是：创建TensorRT推理环境和分配显存空间(TensorRT)–>图像预处理(numpy)–>输入数据复制到显存(TensorRT)–>推理(TensorRT)–>将推理结果复制回CPU端的RAM(TensorRT)–>后处理(numpy)，核心推理代码如下。

1. import tensorrt as trt
2. import pycuda.driver as cuda
3. import pycuda.autoinit  # cuda初始化
4. import numpy as np
6. TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
8. def load_engine(engine_file_path):
9.     """
10.     从engine模型反序列化
11.     :param engine_file_path:
12.     :return:
13.     """
14.     """加载 TensorRT 引擎"""
15.     with open(engine_file_path, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime:
16.         return runtime.deserialize_cuda_engine(f.read())
18. def allocate_buffers(engine):
19.     """
20.     给模型直接输入和输出分配CUDA缓存区
21.     :param engine:
22.     :return:
23.     """
24.     inputs, outputs, bindings = [], [], []  # 记录为输入输出分配的显存空间地址
25.     inputs_shape, outputs_shape = [], []    # 记录从engine模型获取的输入和输出shape
26.     stream = cuda.Stream()
28.     for binding in engine:  # 遍历模型engine中的所有输入输出节点
29.         shape = engine.get_tensor_shape(binding)
30.         size = trt.volume(shape) * np.dtype(np.float32).itemsize
31.         device_mem = cuda.mem_alloc(size)   # 为节点分配size大小的空间，用于储存该节点的数据
32.         engine.get_tensor_mode(binding)
33.         if engine.get_tensor_mode(binding) == trt.TensorIOMode.INPUT:
34.             inputs.append(device_mem)   # 如果是输入节点则将该空间地址存入inputs
35.             inputs_shape.append(shape)  # 记录节点shape
36.         elif engine.get_tensor_mode(binding) == trt.TensorIOMode.OUTPUT:
37.             outputs.append(device_mem)  # 如果是输出节点则将该空间地址存入outputs
38.             outputs_shape.append(shape) # 记录节点shape
39.         bindings.append(int(device_mem))
40.     return inputs, outputs, bindings, stream
43. class YoloTRT:
44.     def __init__(self, engine):
45.         """
46.         构造函数，创建推理环境和分配显存空间
47.         :param engine:
48.         """
49.         self.context = engine.create_execution_context()  # 创建推理的上下文环境
50.         self.inputs, self.outputs, self.bindings, self.stream = allocate_buffers(engine)  # 初始化分配CUDA显存空间
52.     def infer(self, image):
53.         """
54.         推理函数：将RAM（cpu端）的image数据拷贝到GPU，然后进行推理，最后将推理结果数据从GPU上拷贝到RAM（cpu端）中
55.         :param image:
56.         :return:
57.         """
58.         # 将输入数据拷贝到 GPU
59.         cuda.memcpy_htod_async(self.inputs[0], image, self.stream)  # image是cpu的ram数据，需要拷贝到GPU，才能被tensorrt推理使用
60.         # 执行推理
61.         self.context.execute_async_v2(self.bindings, self.stream.handle, None)  # 推理
63.         # 从 GPU 拷贝输出结果到cpu的RAM
64.         # 分配cpu的ram空间，方便从gpu接受输出数据
65.         output0 = np.empty((1, 32, 160, 160), dtype=np.float32)
66.         output1 = np.empty((1, 25200, 117), dtype=np.float32)
67.         cuda.memcpy_dtoh_async(output0, self.outputs[0], self.stream)  # 将位于gpu上的output0输出数据复制到cpu上，方便对齐操作
68.         cuda.memcpy_dtoh_async(output1, self.outputs[1], self.stream)  # 将位于gpu上的output1输出数据复制到cpu上，方便对齐操作
69.         self.stream.synchronize()  # 同步，gpu是并行的，这里相当于等待同步一下
71.         return output0, output1  # 最后将输出结果返回
74. if __name__ == "__main__":
75.     engine_path = "yolov5s-seg.engine"
77.     engine = load_engine(engine_path)
79.     yoloTrt = YoloTRT(engine)  # 初始化一次
81.     # input_data是经过预处理后的输入，shape=(1,3,640,640)，这里用np生成模拟数据
82.     input_data = np.ones((1, 3, 640, 640), np.float32)
83.     output0, output1 = yoloTrt.infer(input_data)  # 推理
84.     print(output0.shape, output1.shape)
85.     # output0.shape= (1, 32, 160, 160)
86.         # output1.shape= (1, 25200, 117)
88.     # 后处理

复制代码

上述代码中没有图像预处理和推理后处理代码，因为这两部门都是用numpy实现，和TensorRT没有关系了，对于认识这类任务的同学已经有丰富的处理履历了，以是只给出TensorRT推理的代码。
对于代码的作用都写了较为详实的备注，如果另有疑惑可以问问chatgpt或deepseek等进行代码详解。
在分配内存时，各变量的详情如下。输入节点有1个，以是分配了一块显存空间。输出节点有两个，以是分配了两块显存空间。

5.2 总结

实际上，无论模子范例（分类、检测、分割等）还是编程语言（Python或C++），TensorRT部署的根本流程通常都遵循以下步调：

• 创建TensorRT推理环境和分配显存空间（TensorRT）
  创建TensorRT的Runtime和Engine对象，分配GPU显存。
  
• 图像预处理
  通常必要将图像数据转换为模子输入所需的格式（如归一化、调整大小等）。在python中可以使用opecv、PIL等库实现，在c++中可以使用opencv或者CUDA编程等方式实现。
  
• 输入数据复制到显存（TensorRT）
  使用TensorRT的API将预处理后的数据从CPU内存复制到GPU显存。在cpu上对图像进行预处理后，必要借助TensorRT的API将数据拷贝到GPU显存，才气让推理直接直接使用到图像数据。
  
• 推理（TensorRT）
  调用TensorRT的enqueue或execute方法进行推理。
  
• 将推理结果复制回CPU端的RAM（TensorRT）
  使用TensorRT的API将推理结果从GPU显存复制回CPU内存。TensorRT推理后，输出数据是直接放在GPU显存上的，python或C++是不能直接利用显存上的数据，以是必要将数据复制回CPU上，才可以进行利用。
  
• 后处理
  通常必要对推理结果进行解码或格式化（如NMS、置信度过滤等）。数据复制回cpu后就和TensorRT无关了，在python或者c++中使用库或者自定义的算法进行后处理即可。
  

总结

• 与TensorRT相关的环节：创建推理环境、分配显存、数据复制、推理执行。
  
• 与TensorRT无关的环节：图像预处理、后处理。
  

这个流程是通用的，实用于各种模子范例和编程语言。具体的实现细节大概会有所不同，但根本步调是同等的。

免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！更多信息从访问主页：qidao123.com:ToB企服之家，中国第一个企服评测及商务社交产业平台。