修改原因
原始yolov8的输出维度是[1, 84, 8400]，本文[1, 5, 8400]会造成后处理的时候会造成内存访问不连续，不利于我们对推理结果做解码。所以修改为[1, 8400, 84](本文[1, 8400, 5])

edit_onnx.py

import onnx
from onnx import helper, shape_inference
from onnx import TensorProto

def modify_onnx_output_shape(onnx_model_path, output_path, new_shape):
    # 加载 ONNX 模型
    model = onnx.load(onnx_model_path)
    graph = model.graph

    # 获取最后一层的输出
    output = graph.output[0]

    # 修改输出的维度
    new_type_proto = helper.make_tensor_type_proto(
        elem_type=output.type.tensor_type.elem_type,  # 保持数据类型不变
        shape=new_shape  # 设置新的输出形状
    )
    output.type.tensor_type.CopyFrom(new_type_proto)

    # 推理新的模型形状（可选，但推荐）
    model = shape_inference.infer_shapes(model)

    # 保存修改后的模型
    onnx.save(model, output_path)
    print(f"Modified model saved to {output_path}")

# 调用函数
modify_onnx_output_shape(
    onnx_model_path="/home/smf/work/Vehicle/ultralytics-main_pp/runs/detect/train11/weights/headout.onnx",
    output_path="/home/smf/work/Vehicle/ultralytics-main_pp/runs/detect/train11/weights/headoutnew.onnx",
    new_shape=[1,8400,5] # [1,8400,48]  # 示例新维度，例如 [batch_size, num_classes]
)

或者修改/ultralytics/nn/modules/head.py

将

return y if self.export else (y, x)

修改为

return y.permute(0, 2, 1) if self.export else (y, x)