将脚本从 pytorch 转换为 onnx 时遇到问题

Question

将 pytorch 脚本转换为 onnx 时遇到问题。我收到错误消息：

onnxruntime.capi.onnxruntime_pybind11_state.InvalidArgument: [ONNXRuntimeError] : 2 : INVALID_ARGUMENT : Unexpected input data type. Actual: (tensor(uint8)) , expected: (tensor(float))

当我查看模型时，它显示：

格式 ONNX v8 制片人 火炬2.1.1 进口 ai.onnx v17 跨步 32 名字 {0: '玩家'} 图形 主图 图片 名称： 图片 张量：float32[1,3,320,320] 输出0 名称：输出0 张量：float32[1,6300,6]

据我了解是 float32

那么为什么说数据类型是uint8？

#model export command: python .\export.py --weights ./best.pt --include onnx --imgsz 320 320 --device 0

加载ONNX模型

   ` print("[INFO] Loading ONNX model")
    self.onnx_session = onnxruntime.InferenceSession(onnx_model_path)`

我试过了

确保输入张量具有正确的形状

`frame = frame.transpose((2, 0, 1))  # Change from HWC to CHW format
frame = frame[np.newaxis, :]  # Add batch size dimension
frame = frame.astype(np.float32) / 255.0
results = self.onnx_session.run(['output0'], {'images': frame})`

该代码通过将图像帧的格式更改为 CHW（通道、高度、宽度）、添加批量大小维度并标准化像素值，为 ONNX 模型推理准备图像帧。然后将处理后的图像输入 ONNX 模型进行推理，并将结果存储在 results 变量中。

这导致了这个

onnxruntime.capi.onnxruntime_pybind11_state.InvalidArgument: [ONNXRuntimeError] : 2 : INVALID_ARGUMENT : Got invalid dimensions for input: images for the following indices index: 1 Got: 4 Expected: 3 Please fix either the inputs or the model.

我错过了什么？

Answer 1

从您收到的错误消息中

Got invalid dimensions for input: images for the following indices index: 1
Got: 4 Expected: 3

我们看到推理失败，因为 ONNX 模型期望输入数组在第二维具有三个通道，而实际输入有四个。考虑到您的输入是图像，

images tensor: float32[1,3,320,320] 

第二个维度最有可能用于红色、绿色和蓝色通道。由于您的输入有一个通道太多，我敢打赌您的图像采用

RGBA

格式，因此会出现一个额外的通道 (alpha) 和错误消息。

在将图像输入图表之前，尝试将其转换为 RGB。

from PIL import Image
image = Image.open("/path/to/image")
image = image.convert("RGB")
frame = np.asarray(image)
...