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模子：
python export.py --weights=yolov5s-seg.pt --imgsz= --include==onnx
通过以上命令就可以导出onnx模子。export.py还可以导出更多范例，并定制一些需求，可以浏览export.py学习更多。
2.2 测试
起首，我们通过netron可视化onnx模子。可以看到该分割模子的输入为，有两个输出：output0和output1
https://i-blog.csdnimg.cn/direct/1748bf41c2a44fafa6a4f239755daddd.png
为了方便后续进行后处理（虽然和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模子。在这个过程中会自动根据装备的特异性进行优化，该过程十分缓慢。但是完成后会生存为本地文件，下次可以直接加载使用，不用再次序列化。
https://i-blog.csdnimg.cn/direct/83478dd532324ad2a8dcca9d37c9e38a.png
这个过程可以使用基于python的API完成，或者直接使用trtexec工具。
4.1 使用trtexec工具完成序列化
在第3节中下载了TensorRT包，在bin中有trtexec工具。打开终端进入trtexec工具所在文件夹就可以使用该命令工具了。
https://i-blog.csdnimg.cn/direct/4a7885eedbef475b996de1bbe2527798.png
转换命令为：
trtexec --onnx=yolov5s-seg.onnx --saveEngine=yolov5s-seg.engine --fp16
onnx指定onnx模子路径，saveEngine知道转换后engine模生存路径，fp16表现使用FP16，这个可以大大提高推理速度。不加–fp16则默认使用FLOAT32
末了会出现PASSED结果，说明成功了，此时就得到engine模子。如下图所示，还会给出一些推理时间、内存占用等测试信息。
https://i-blog.csdnimg.cn/direct/2c8bfbc2cbef44c98f41fedf2dcfaab2.png
更多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 模子：
trtexec --onnx=model.onnx --saveEngine=model.engine
[*] 性能测试：
trtexec --loadEngine=model.engine --batch=16 --fp16
[*] 使用 INT8 精度：
trtexec --onnx=model.onnx --int8 --batch=32
总结
trtexec 是 TensorRT 的重要工具，支持模子转换、性能测试和推理执行，通过调整参数可以优化模子性能。
4.2 使用python的API进行转换
除了使用trtexec工具进行转换，也可以使用python的API进行转换，但是二者没有什么区别。
使用python的API进行转换具有比力固定的流程，只不过可以根据需求比如动态输入等进行相应的设置。
不管什么方式转换得到的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推理环境和分配显存空间(TensorRT)–>图像预处理(numpy)–>输入数据复制到显存(TensorRT)–>推理(TensorRT)–>将推理结果复制回CPU端的RAM(TensorRT)–>后处理(numpy)，核心推理代码如下。
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit# cuda初始化
import numpy as np
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
def load_engine(engine_file_path):
"""
从engine模型反序列化
:param engine_file_path:
:return:
"""
"""加载 TensorRT 引擎"""
with open(engine_file_path, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime:
return runtime.deserialize_cuda_engine(f.read())
def allocate_buffers(engine):
"""
给模型直接输入和输出分配CUDA缓存区
:param engine:
:return:
"""
inputs, outputs, bindings = [], [], []# 记录为输入输出分配的显存空间地址
inputs_shape, outputs_shape = [], [] # 记录从engine模型获取的输入和输出shape
stream = cuda.Stream()
for binding in engine:# 遍历模型engine中的所有输入输出节点
shape = engine.get_tensor_shape(binding)
size = trt.volume(shape) * np.dtype(np.float32).itemsize
device_mem = cuda.mem_alloc(size) # 为节点分配size大小的空间，用于储存该节点的数据
engine.get_tensor_mode(binding)
if engine.get_tensor_mode(binding) == trt.TensorIOMode.INPUT:
inputs.append(device_mem) # 如果是输入节点则将该空间地址存入inputs
inputs_shape.append(shape)# 记录节点shape
elif engine.get_tensor_mode(binding) == trt.TensorIOMode.OUTPUT:
outputs.append(device_mem)# 如果是输出节点则将该空间地址存入outputs
outputs_shape.append(shape) # 记录节点shape
bindings.append(int(device_mem))
return inputs, outputs, bindings, stream
class YoloTRT:
def __init__(self, engine):
"""
构造函数，创建推理环境和分配显存空间
:param engine:
"""
self.context = engine.create_execution_context()# 创建推理的上下文环境
self.inputs, self.outputs, self.bindings, self.stream = allocate_buffers(engine)# 初始化分配CUDA显存空间
def infer(self, image):
"""
推理函数：将RAM（cpu端）的image数据拷贝到GPU，然后进行推理，最后将推理结果数据从GPU上拷贝到RAM（cpu端）中
:param image:
:return:
"""
# 将输入数据拷贝到 GPU
cuda.memcpy_htod_async(self.inputs, image, self.stream)# image是cpu的ram数据，需要拷贝到GPU，才能被tensorrt推理使用
# 执行推理
self.context.execute_async_v2(self.bindings, self.stream.handle, None)# 推理
# 从 GPU 拷贝输出结果到cpu的RAM
# 分配cpu的ram空间，方便从gpu接受输出数据
output0 = np.empty((1, 32, 160, 160), dtype=np.float32)
output1 = np.empty((1, 25200, 117), dtype=np.float32)
cuda.memcpy_dtoh_async(output0, self.outputs, self.stream)# 将位于gpu上的output0输出数据复制到cpu上，方便对齐操作
cuda.memcpy_dtoh_async(output1, self.outputs, self.stream)# 将位于gpu上的output1输出数据复制到cpu上，方便对齐操作
self.stream.synchronize()# 同步，gpu是并行的，这里相当于等待同步一下
return output0, output1# 最后将输出结果返回
if __name__ == "__main__":
engine_path = "yolov5s-seg.engine"
engine = load_engine(engine_path)
yoloTrt = YoloTRT(engine)# 初始化一次
# input_data是经过预处理后的输入，shape=(1,3,640,640)，这里用np生成模拟数据
input_data = np.ones((1, 3, 640, 640), np.float32)
output0, output1 = yoloTrt.infer(input_data)# 推理
print(output0.shape, output1.shape)
# output0.shape= (1, 32, 160, 160)
# output1.shape= (1, 25200, 117)
# 后处理
上述代码中没有图像预处理和推理后处理代码，因为这两部门都是用numpy实现，和TensorRT没有关系了，对于认识这类任务的同学已经有丰富的处理履历了，以是只给出TensorRT推理的代码。
对于代码的作用都写了较为详实的备注，如果另有疑惑可以问问chatgpt或deepseek等进行代码详解。
在分配内存时，各变量的详情如下。输入节点有1个，以是分配了一块显存空间。输出节点有两个，以是分配了两块显存空间。
https://i-blog.csdnimg.cn/direct/2d8c7cbfd0af4ca1a6b875d8c8f46333.png
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企服之家，中国第一个企服评测及商务社交产业平台。
页:
[1]