【找出 ONNX 输入】onnx_model.graph.input 或者 ort_session.get_inputs()

先下载好 onnx 和 onnxruntime-gpu

pip install onnx
pip install onnxruntime-gpu

这里需要注意下 onnxruntime 对应的 CUDA 和 cuDNN 版本，不一定适配，具体可查看官方文档

针对转化成果的 .onnx 模型，通过以下代码即可查看定义好的输入和输出

import onnx
import onnxruntime as ort

# 加载 ONNX 模型
onnx_model_path = "/path/to/yourmodel.onnx"
onnx_model = onnx.load(onnx_model_path)

# 打印模型输入定义列表
print("ONNX 模型的输入定义:")
for input in onnx_model.graph.input:
    print(f"Name: {input.name}")
    print(f"Type: {input.type}")
    print(f"Shape: {[dim.dim_value for dim in input.type.tensor_type.shape.dim]}")

# 使用 onnxruntime 加载模型并打印输入定义
ort_session = ort.InferenceSession(onnx_model_path)
print("\nONNX Runtime 模型的输入定义:")
for input in ort_session.get_inputs():
    print(f"Name: {input.name}")
    print(f"Type: {input.type}")
    print(f"Shape: {input.shape}")

正确的输出情况示例如下，其中为 0 的维度是可变的维度

ONNX 模型的输入定义:
Name: input
Type: tensor_type {
  elem_type: 1
  shape {
    dim {
      dim_param: "batch_size"
    }
    dim {
      dim_value: 3
    }
    dim {
      dim_param: "height"
    }
    dim {
      dim_param: "width"
    }
  }
}

Shape: [0, 3, 0, 0]

ONNX Runtime 模型的输入定义:
Name: input
Type: tensor(float)
Shape: ['batch_size', 3, 'height', 'width']

模型有问题时的输出结果如下：

ONNX 模型的输入定义:
2024-05-22 17:08:34.460766609 [W:onnxruntime:, constant_folding.cc:269 ApplyImpl] Could not find a CPU kernel and hence can't constant fold Add node '/Add'
2024-05-22 17:08:34.461292057 [W:onnxruntime:, constant_folding.cc:269 ApplyImpl] Could not find a CPU kernel and hence can't constant fold Sqrt node '/transformer/resblocks/resblocks.0/attn/Sqrt'
2024-05-22 17:08:34.461538350 [W:onnxruntime:, constant_folding.cc:269 ApplyImpl] Could not find a CPU kernel and hence can't constant fold Sqrt node '/transformer/resblocks/resblocks.1/attn/Sqrt'
2024-05-22 17:08:34.461734198 [W:onnxruntime:, constant_folding.cc:269 ApplyImpl] Could not find a CPU kernel and hence can't constant fold Sqrt node '/transformer/resblocks/resblocks.2/attn/Sqrt'
2024-05-22 17:08:34.461935176 [W:onnxruntime:, constant_folding.cc:269 ApplyImpl] Could not find a CPU kernel and hence can't constant fold Sqrt node '/transformer/resblocks/resblocks.3/attn/Sqrt'
2024-05-22 17:08:34.462130423 [W:onnxruntime:, constant_folding.cc:269 ApplyImpl] Could not find a CPU kernel and hence can't constant fold Sqrt node '/transformer/resblocks/resblocks.4/attn/Sqrt'
2024-05-22 17:08:34.462322935 [W:onnxruntime:, constant_folding.cc:269 ApplyImpl] Could not find a CPU kernel and hence can't constant fold Sqrt node '/transformer/resblocks/resblocks.5/attn/Sqrt'
2024-05-22 17:08:34.462515127 [W:onnxruntime:, constant_folding.cc:269 ApplyImpl] Could not find a CPU kernel and hence can't constant fold Sqrt node '/transformer/resblocks/resblocks.6/attn/Sqrt'
2024-05-22 17:08:34.462715684 [W:onnxruntime:, constant_folding.cc:269 ApplyImpl] Could not find a CPU kernel and hence can't constant fold Sqrt node '/transformer/resblocks/resblocks.7/attn/Sqrt'
2024-05-22 17:08:34.462897496 [W:onnxruntime:, constant_folding.cc:269 ApplyImpl] Could not find a CPU kernel and hence can't constant fold Sqrt node '/transformer/resblocks/resblocks.8/attn/Sqrt'
2024-05-22 17:08:34.463080460 [W:onnxruntime:, constant_folding.cc:269 ApplyImpl] Could not find a CPU kernel and hence can't constant fold Sqrt node '/transformer/resblocks/resblocks.9/attn/Sqrt'
2024-05-22 17:08:34.463272601 [W:onnxruntime:, constant_folding.cc:269 ApplyImpl] Could not find a CPU kernel and hence can't constant fold Sqrt node '/transformer/resblocks/resblocks.10/attn/Sqrt'
2024-05-22 17:08:34.463454613 [W:onnxruntime:, constant_folding.cc:269 ApplyImpl] Could not find a CPU kernel and hence can't constant fold Sqrt node '/transformer/resblocks/resblocks.11/attn/Sqrt'
2024-05-22 17:08:34.478391998 [W:onnxruntime:, constant_folding.cc:269 ApplyImpl] Could not find a CPU kernel and hence can't constant fold Sqrt node '/transformer/resblocks/resblocks.11/attn/Sqrt'
2024-05-22 17:08:34.478412346 [W:onnxruntime:, constant_folding.cc:269 ApplyImpl] Could not find a CPU kernel and hence can't constant fold Sqrt node '/transformer/resblocks/resblocks.10/attn/Sqrt'
2024-05-22 17:08:34.478425841 [W:onnxruntime:, constant_folding.cc:269 ApplyImpl] Could not find a CPU kernel and hence can't constant fold Sqrt node '/transformer/resblocks/resblocks.9/attn/Sqrt'
2024-05-22 17:08:34.478437954 [W:onnxruntime:, constant_folding.cc:269 ApplyImpl] Could not find a CPU kernel and hence can't constant fold Sqrt node '/transformer/resblocks/resblocks.8/attn/Sqrt'
2024-05-22 17:08:34.478449596 [W:onnxruntime:, constant_folding.cc:269 ApplyImpl] Could not find a CPU kernel and hence can't constant fold Sqrt node '/transformer/resblocks/resblocks.7/attn/Sqrt'
2024-05-22 17:08:34.478461719 [W:onnxruntime:, constant_folding.cc:269 ApplyImpl] Could not find a CPU kernel and hence can't constant fold Sqrt node '/transformer/resblocks/resblocks.6/attn/Sqrt'
2024-05-22 17:08:34.478473241 [W:onnxruntime:, constant_folding.cc:269 ApplyImpl] Could not find a CPU kernel and hence can't constant fold Sqrt node '/transformer/resblocks/resblocks.5/attn/Sqrt'
2024-05-22 17:08:34.478484973 [W:onnxruntime:, constant_folding.cc:269 ApplyImpl] Could not find a CPU kernel and hence can't constant fold Sqrt node '/transformer/resblocks/resblocks.4/attn/Sqrt'
2024-05-22 17:08:34.478496114 [W:onnxruntime:, constant_folding.cc:269 ApplyImpl] Could not find a CPU kernel and hence can't constant fold Sqrt node '/transformer/resblocks/resblocks.3/attn/Sqrt'
2024-05-22 17:08:34.478507375 [W:onnxruntime:, constant_folding.cc:269 ApplyImpl] Could not find a CPU kernel and hence can't constant fold Sqrt node '/transformer/resblocks/resblocks.2/attn/Sqrt'
2024-05-22 17:08:34.478518706 [W:onnxruntime:, constant_folding.cc:269 ApplyImpl] Could not find a CPU kernel and hence can't constant fold Sqrt node '/transformer/resblocks/resblocks.1/attn/Sqrt'
2024-05-22 17:08:34.478529707 [W:onnxruntime:, constant_folding.cc:269 ApplyImpl] Could not find a CPU kernel and hence can't constant fold Sqrt node '/transformer/resblocks/resblocks.0/attn/Sqrt'
2024-05-22 17:08:34.479066035 [W:onnxruntime:, constant_folding.cc:269 ApplyImpl] Could not find a CPU kernel and hence can't constant fold Add node '/Add'

ONNX Runtime 模型的输入定义: