将 ONNX 模型的输入布局从 NCHW 更改为 NHWC

Question

我有一个输入形状为

(1, 1, 28, 28)

的 ONNX 模型，它采用

NCHW

布局。 有什么方法可以将模型转换为在

NHWC

布局中获取输入数据，即

(1, 28, 28, 1)

？

使用以下代码从文件加载模型：

model_obj = onnx.load(model_path)

修改和转换 ONNX 文件的方法也很有帮助。

Answer 1

ONNX 仅支持 NCHW，但是您可以修改模型以采用 NHWC，然后添加转置。

我没有示例可以向您展示，但是我有一些您可能会觉得有用的资源。

1. 在 GPU 上运行 ONNX 的 ProViz 示例中，有一个 CUDA 内核的示例，该内核执行 NHWC 到 NCHW（反之亦然）https://github.com/NVIDIA/ProViz-AI-Samples/blob /master/onnxruntime_cpp_samples/cuda_provider/src/FormatConversion.cu

2. ProViz 关于上述存储库的博文（他们简要提到了转换以及如何处理它）：https://developer.nvidia.com/blog/end-to-end-ai-for-nvidia-based- pcs-cuda-and-tensorrt-execution-providers-in-onnx-runtime/#nhwc_format

很抱歉没有给出最有用的答案，但这就是我所知道的有关该主题的全部内容。