使用Tensorrt的python api 部署支持动态batch的yolov5s

最新推荐文章于 2025-09-02 13:17:01 发布

原创

最新推荐文章于 2025-09-02 13:17:01 发布 · 5.1k 阅读

30 ·

CC 4.0 BY-SA版权

文章标签：

#机器学习 #tensorrt

本文档详细介绍了如何在Ubuntu 18.04上使用TensorRT 8.2和PyTorch 1.10.0构建YOLOv5s的ONNX模型，并通过动态batch支持生成TensorRT引擎。首先，修改YOLOv5源码简化模型输出，然后导出ONNX模型。接着，使用trtexec的替代方法，创建包含优化配置文件的TensorRT引擎，处理动态维度。最后，演示了如何从引擎文件反序列化并进行推理，包括数据预处理、推理和结果处理。

系统环境:

ubuntu 1804

cuda11.3

tensorrt 8.2.06

显卡2080

pytorch 1.10.0

onnx 1.10.2

onnx-simplifier 0.3.6

步骤１：导出onnx 模型(参考https://github.com/shouxieai/tensorRT_Pro)

1.1 下载yolov5项目并修改相关代码，目的是减少导出onnx的复杂度，只保留一个输出便于后处理

# 下载并进入yolov5项目
git clone git@github.com:ultralytics/yolov5.git
cd yolov5/models
gedit yolo.py

#更改yolov5/models/yolo.py Detect.forward函数 

# yolov5/models/yolo.py 
# bs, _, ny, nx = x[i].shape  # x(bs,255,20,20) to x(bs,3,20,20,85)
# x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
# 修改为:
bs, _, ny, nx = x[i].shape  # x(bs,255,20,20) to x(bs,3,20,20,85)
bs = -1
ny = int(ny)
nx = int(nx)
x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()

# yolov5/models/yolo.py
#  z.append(y.view(bs, -1, self.no))
# 修改为：
z.append(y.view(bs, self.na * ny * nx, self.no))

############# 对于 yolov5-6.0 #####################
# yolov5/models/yolo.py Detect.forward函数
# if self.grid[i].shape[2:4] != x[i].shape[2:4] or self.onnx_dynamic:
#    self.grid[i], self.anchor_grid[i] = self._make_grid(nx, ny, i)
# 修改为:
if self.grid[i].shape[2:4] != x[i].shape[2:4] or self.onnx_dynamic:
    self.grid[i], self.anchor_grid[i] = self._make_grid(nx, ny, i)

    # disconnect for pytorch trace
    anchor_grid = (self.anchors[i].clone() * self.stride[i]).view(1, -1, 1, 1, 2)

# y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i]  # wh
# 修改为:
y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * anchor_grid  # wh


# wh = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i]  # wh
# 修改为:
wh = (y[..., 2:4] * 2) ** 2 * anchor_grid  # wh

#return x if self.training else (torch.cat(z, 1), x)
#修改为
return torch.cat(z, 1)

1.2 模型导出

修改yolov5/export.py,目的是只保留onnx模型动态batch

############# 对于 yolov5-6.0 #####################

# yolov5/export.py　export_onnx函数
# torch.onnx.export(model, im, f, verbose=False, opset_version=opset,
#                   training=torch.onn

最低0.47元/天解锁文章