yolov5 pt转engine的问题

Question

帮忙解决yolov5的pt转engine的问题。
问题是：我标注图片设为1个分类时用yolov5训练，然后将pt转成onnx，再转engine运行都没问题。
转onnx我用的是yolov5原版的export.py，只改了图像尺寸（imgsz）为2400，其它基本没改。
转engine是别人给的一段代码。
当我标注分类为4个类别时，将pt转成onnx没问题，但转engine时报如下错误：

[1,3,300,300,3] and [1,3,300,300,2]).
LOG[1]: /model.24/Add: elementwise inputs must have same dimensions or follow broadcast rules (input dimensions were [1,3,300,300,3] and [1,3,300,300,2]).
LOG[1]: /model.24/Add: elementwise inputs must have same dimensions or follow broadcast rules (input dimensions were [1,3,300,300,3] and [1,3,300,300,2]).
ERROR: onnx2trt_utils.cpp:680 In function elementwiseHelper:
[8] Assertion failed: tensor_ptr->getDimensions().nbDims == maxNbDims && "Failed to broadcast tensors elementwise!"
Parse failed.

当我标注分类为2个类别时，将pt转成onnx没问题，但转engine时报如下错误：
LOG[1]: /model.24/Reshape_1: volume mismatch. Input dimensions [1,3,300,300,6] have volume 1620000 and output dimensions [1,270000,7] have volume 1890000.
ERROR: onnx2trt_utils.cpp:188 In function convertAxis:
[8] Assertion failed: axis >= 0 && axis < nbDims
Parse failed.

export.py中代码如下：

def export_onnx(model, im, file, opset, dynamic, simplify, prefix=colorstr('ONNX:')):
    # YOLOv5 ONNX export
    check_requirements('onnx>=1.12.0')
    import onnx

    LOGGER.info(f'\n{prefix} starting export with onnx {onnx.__version__}...')
    f = file.with_suffix('.onnx')

    output_names = ['output0', 'output1'] if isinstance(model, SegmentationModel) else ['output0']
    if dynamic:
        dynamic = {'images': {0: 'batch', 2: 'height', 3: 'width'}}  # shape(1,3,640,640)
        if isinstance(model, SegmentationModel):
            dynamic['output0'] = {0: 'batch', 1: 'anchors'}  # shape(1,25200,85)
            dynamic['output1'] = {0: 'batch', 2: 'mask_height', 3: 'mask_width'}  # shape(1,32,160,160)
        elif isinstance(model, DetectionModel):
            dynamic['output0'] = {0: 'batch', 1: 'anchors'}  # shape(1,25200,85)

    torch.onnx.export(
        model.cpu() if dynamic else model,  # --dynamic only compatible with cpu
        im.cpu() if dynamic else im,
        f,
        verbose=False,
        opset_version=opset,
        do_constant_folding=True,  # WARNING: DNN inference with torch>=1.12 may require do_constant_folding=False
        input_names=['images'],
        output_names=output_names,
        dynamic_axes=dynamic or None)

    # Checks
    model_onnx = onnx.load(f)  # load onnx model
    onnx.checker.check_model(model_onnx)  # check onnx model

    # Metadata
    d = {'stride': int(max(model.stride)), 'names': model.names}
    for k, v in d.items():
        meta = model_onnx.metadata_props.add()
        meta.key, meta.value = k, str(v)
    onnx.save(model_onnx, f)

    # Simplify
    if simplify:
        try:
            cuda = torch.cuda.is_available()
            check_requirements(('onnxruntime-gpu' if cuda else 'onnxruntime', 'onnx-simplifier>=0.4.1'))
            import onnxsim

            LOGGER.info(f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...')
            model_onnx, check = onnxsim.simplify(model_onnx)
            assert check, 'assert check failed'
            onnx.save(model_onnx, f)
        except Exception as e:
            LOGGER.info(f'{prefix} simplifier failure: {e}')
    return f, model_onnx

网上查找说是需要变换写法，但我找不到对应修改位置。
哪位帮看看怎么解决，最好有具体修改方法，我是初学者，只会简单的操作。

Answer 1

结合GPT给出回答如下请题主参考
首先，将 PyTorch 模型转换为 ONNX 格式:

import torch
from torch.autograd import Variable
import torchvision
import onnx
import onnxruntime

# Load the model
model = torch.load('path/to/your/pytorch/model.pth')

# Set the model to inference mode
model.eval()

# Export the model to ONNX format
dummy_input = Variable(torch.randn(1, 3, width, height)) # Set input image dimensions
input_names = ['image']
output_names = ['boxes', 'scores', 'classes'] # Assuming you have 1 class
torch.onnx.export(model, dummy_input, 'path/to/output/onnx/model.onnx', verbose=True, input_names=input_names, output_names=output_names)

然后，将 ONNX 模型转换为 TensorRT engine：

import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np

TRT_LOGGER = trt.Logger(trt.Logger.WARNING)

def build_engine(onnx_file_path, engine_file_path="path/to/output/tensorrt/engine/tensorrt.engine"):
    """Takes an ONNX file and creates a TensorRT engine to run inference with"""
    with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network, trt.OnnxParser(network, TRT_LOGGER) as parser:
        builder.max_workspace_size = 1 << 30 # 1GB
        builder.max_batch_size = 1
        # Parse the ONNX file to populate the TensorRT network
        with open(onnx_file_path, 'rb') as model:
            parser.parse(model.read())
        # Generate the TensorRT engine optimized for the target platform
        with builder.build_cuda_engine(network) as engine:
            with open(engine_file_path, "wb") as f:
                f.write(engine.serialize())
            return engine

def allocate_buffers(engine):
    """Allocates memory for input and output buffers on the device"""
    inputs = []
    outputs = []
    bindings = []
    stream = cuda.Stream()
    for binding in engine:
        size = trt.volume(engine.get_binding_shape(binding)) * engine.max_batch_size
        dtype = trt.nptype(engine.get_binding_dtype(binding))
        # Allocate device memory
        buf = cuda.mem_alloc(size * dtype.itemsize)
        bindings.append(int(buf))
        # Append to the appropriate list
        if engine.binding_is_input(binding):
            inputs.append(buf)
        else:
            outputs.append(buf)
    return inputs, outputs, bindings, stream

def do_inference(engine, inputs, outputs, bindings, stream):
    """Runs inference on a TensorRT engine"""
    # Transfer input data to the GPU
    cuda.memcpy_htod_async(inputs[0], np.ascontiguousarray(image), stream)
    # Run inference
    engine.execute_async_v2(bindings=bindings, stream_handle=stream.handle)
    # Transfer predictions back from the GPU
    cuda.memcpy_dtoh_async(outputs[0], output_data, stream)
    # Synchronize the stream
    stream.synchronize()
    # Return the output predictions
    return output_data

最后，您可以加载您的 TensorRT engine 并执行推理：

# Load the TensorRT engine
engine = trt.lite.Engine.deserialize(engine_file_path)

# Allocate memory for input and output buffers
inputs, outputs, bindings, stream = allocate_buffers(engine)

# Load the input image
image = cv2.imread("path/to/your/input/image.jpg")

# Preprocess the input image (e.g. resize, normalize, etc.)

# Run inference
output_data = do_inference(engine, inputs, outputs, bindings, stream)

# Postprocess the output data (e.g. decode bounding boxes, apply NMS, etc.)

希望以上代码可以帮助您将 YOLOv5 的 PyTorch 模型转换为 TensorRT engine，并执行推理。