YOLO模型转ONNX后在ONNX Runtime运行返回结果形状不一致问题

• 你所用的是YOLOv3架构的检测模型，原生设计就包含3个不同下采样倍率的输出头，分别对应大、中、小不同尺度的目标检测，因此原始Keras模型推理时返回3个张量组成的列表属于正常情况。
• 出现输出数量不一致的核心原因有两类：
  1. 推理ONNX模型时仅提取了第一个输出节点，转换后的模型实际仍保留了全部3个输出
  2. SavedModel导出或tf2onnx转换阶段没有显式指定多输出配置，导致输出被自动合并/截断

第一步：验证ONNX模型的实际输出数量 #

先运行如下代码查看转换后ONNX模型的所有输出节点：

import onnxruntime as rt
sess = rt.InferenceSession("drive/MyDrive/ONNX_demo/models/model.onnx")
# 遍历打印所有输出节点信息
for idx, output in enumerate(sess.get_outputs()):
    print(f"输出序号{idx}：名称={output.name}，预期形状={output.shape}")

如果打印结果显示存在3个输出，直接修改推理代码，提取全部输出即可：

input_name = sess.get_inputs()[0].name
output_names = [out.name for out in sess.get_outputs()]
yolo_pred = sess.run(output_names, {input_name: img.astype(np.float32)})

此时得到的yolo_pred就和原始模型输出结构完全一致。

第二步：修正导出&转换配置（如果ONNX确实只有1个输出） #

如果上一步查询到ONNX仅存在1个输出，需要按如下流程重新导出和转换模型：

1. 导出SavedModel时显式指定多输出签名，避免输出被自动裁剪：

import tensorflow as tf
# 构建签名，显式指定所有输入输出
signatures = tf.saved_model.signature_def_utils.predict_signature_def(
    inputs={"input": model.input},
    outputs={
        "output_0": model.output[0],
        "output_1": model.output[1],
        "output_2": model.output[2]
    }
)
tf.saved_model.save(
    model,
    "drive/MyDrive/ONNX_demo/models/save_model_fix",
    signatures=signatures
)

2. 转换时显式指定所有输出节点，禁止tf2onnx自动合并输出：
   首先用命令行查询SavedModel的输出节点名：

saved_model_cli show --dir drive/MyDrive/ONNX_demo/models/save_model_fix --all

在输出的signature_def部分找到3个输出对应的节点名（通常默认是Identity、Identity_1、Identity_2），替换到如下转换命令中：

python3 -m tf2onnx.convert \
--saved-model "drive/MyDrive/ONNX_demo/models/save_model_fix" \
--output "drive/MyDrive/ONNX_demo/models/model_fix.onnx" \
--outputs Identity,Identity_1,Identity_2

转换完成后再按第一步的方式提取全部3个输出即可匹配原始模型的推理结果。

内容的提问来源于stack exchange，提问作者Luis Ramon Ramirez Rodriguez