yolov8/yolov11 bb如何转出多 batch 的模型, 以及修改onnx模型输出维度.
官方的yolov8/yolov11 obb的转 onnx, 它输出的维度是 :

(1, x, y, w, h, cls_conf, angle, anchors)

但是这样的维度不⽅便用Tensorrt进行本地代码部署和解析。需要将onnx 的输出维度改为：

(1, anchors, x, y, w, h, cls_conf, angle)

那么, 需要修改代码, 把修改输入和输出节点名和多batch推理。

一. 两处代码修改:

修改一:  修改输入和输出节点名, 以及多batch:

(ultralytics-main/ultralytics/engine/exporter.py)

在class Exporter:这个类中def export_onnx(self, prefix=colorstr("ONNX:")):

(找到对应的代码, 大约在409行左右, 注释这段代码, 添加修改后的代码)

        # output_names = ["output0", "output1"] if isinstance(self.model, SegmentationModel) else ["output0"]
        # dynamic = self.args.dynamic
        # if dynamic:
        #     dynamic = {"images": {0: "batch", 2: "height", 3: "width"}}  # shape(1,3,640,640)
        #     if isinstance(self.model, SegmentationModel):
        #         dynamic["output0"] = {0: "batch", 2: "anchors"}  # shape(1, 116, 8400)
        #         dynamic["output1"] = {0: "batch", 2: "mask_height", 3: "mask_width"}  # shape(1,32,160,160)
        #     elif isinstance(self.model, DetectionModel):
        #         dynamic["output0"] = {0: "batch", 2: "anchors"}  # shape(1, 84, 8400)
        #


        # (修改的代码)转模型修改为下面
        output_names = ['output0', 'output1'] if isinstance(self.model, SegmentationModel) else ['output']
        dynamic = self.args.dynamic
        if dynamic:
            dynamic = {'images': {0: 'batch'}}  # shape(1,3,640,640)
            if isinstance(self.model, SegmentationModel):
                dynamic['output0'] = {0: 'batch', 2: 'anchors'}  # shape(1, 116, 8400)
                dynamic['output1'] = {0: 'batch', 2: 'mask_height', 3: 'mask_width'}  # shape(1,32,160,160)
            elif isinstance(self.model, DetectionModel):
                dynamic['output'] = {0: 'batch'}  # shape(1, 84, 8400)

修改二: 修改模型输出维度

需要将官方代码的输出维度[1, 类别数+5, 8400], 修改为[1, 8400, 类别数+5] , 我们的tensorrt代码需要这个维度, 在做concat 的时候就做维度修改.

ultralytics-main/ultralytics/nn/modules/head.py

在类class OBB(Detect)中的def forward(self, x):

# 修改最后两个维度, 维度从[1, 类别数+5, 8400], 修改为[1, 8400, 类别数+5]
return torch.cat([x, angle], 1).permute(0, 2, 1) if self.export else (torch.cat([x[0], angle], 1), (x[1], angle))

(注: 模型训练和推理, 不需要把exporter.py和head.py的代码改回去, 修改后的代码能正常训练和推理)

二.pt转onnx模型

from ultralytics import YOLO

# Load a model

model = YOLO('xxx/weights/best.pt')

# 转化为256 *640的宽高,这里高在前，宽在后，
# dynamic=True 表示导出的ONNX模型将支持动态输入尺寸，允许你在推理时使用不同尺寸的图像。
success = model.export(format="onnx", dynamic=False, simplify=True, imgsz=(640, 640))

三. onnx转engine（多batch推理）

由于要用到trtexec，而trtexec在/usr/src/tensorrt/bin目录下，每次执行，需要cd到该目录下执行下面的代码，这样有点麻烦，可以创建软链接放到要执行指令的目录下即可。比如我软链接到了Download目录下.

ln -n /usr/src/tensorrt/bin/trtexec ./

onnx在jetson nx板子上转engine:

检测模型转多batch推理的engine:

# 检测模型转换batch推理的engine
（注意：--minShapes：1， --optShapes：2， --maxShapes：4）

cd Download
# 修改下面参数--fp16转为fp16,  --int8转为int8
# TensorRT ≤ 8.4:
sudo ./trtexec --onnx=yolo8s-obb_640_640.onnx --saveEngine=yolo8s-obb_640_640.engine --workspace=4096 --fp16 --minShapes=images:1x3x640x640 --optShapes=images:2x3x640x640 --maxShapes=images:4x3x640x640

# TensorRT ≥ 8.5
sudo ./trtexec --onnx=yolo8s-obb_640_640.onnx --saveEngine=yolo8s-obb_640_640.engine --memPoolSize=workspace:4096 --fp16 --minShapes=images:1x3x640x640 --optShapes=images:2x3x640x640 --maxShapes=images:4x3x640x640

分类模型转多batch推理的engine: 

# 分类模型转换命令(多batch)：
（注意：--minShapes：4， --optShapes：4， --maxShapes：4）
# 修改下面参数--fp16转为fp16,  --int8转为int8
# TensorRT ≤ 8.4
sudo ./trtexec --onnx=yolov8s_68_labels_8_24_.onnx --saveEngine=yolov8s_68_labels_8_24_.engine --fp16 --minShapes=images:4x3x224x224 --optShapes=images:4x3x224x224 --maxShapes=images:4x3x224x224

到这里, 转多batch推理的engine模型就完成了.