ONNX 原生仅支持 NCHW 格式，但可通过修改计算图——在输入后插入 Transpose 节点，使模型逻辑上接受 NHWC 输入（如 (1, 28, 28, 1)），并自动转为内部所需的 NCHW 格式。本文提供完整的 Python 实现方案，包括图结构修改、形状更新与模型验证。
onnx 原生仅支持 nchw 格式，但可通过修改计算图——在输入后插入 transpose 节点，使模型逻辑上接受 nhwc 输入（如 `(1, 28, 28, 1)`），并自动转为内部所需的 nchw 格式。本文提供完整的 python 实现方案，包括图结构修改、形状更新与模型验证。
要让原本设计为 NCHW（[batch, channel, height, width]）输入的 ONNX 模型兼容 NHWC（[batch, height, width, channel]）格式，不能直接更改 ONNX 的“原生布局约定”（ONNX 规范中卷积等算子默认以 NCHW 解释张量），但可通过在模型前端插入显式 Transpose 节点实现语义转换：模型对外暴露 NHWC 输入接口，内部立即转为 NCHW 继续执行。这是一种标准、安全且可导出的图改写方式。
以下是使用 onnx Python 库完成该转换的完整步骤：
✅ 步骤 1：加载并解析原始模型
import onnx
from onnx import helper, TensorProto
model = onnx.load("model.onnx")
graph = model.graph
✅ 步骤 2：修改输入定义（形状 + 名称）
假设原始输入名为 "input"，形状为 (1, 1, 28, 28)（NCHW）：
# 获取原始输入
original_input = graph.input[0]
# 创建新的 NHWC 输入：(1, 28, 28, 1)
new_input = helper.make_tensor_value_info(
"input_nhwc",
original_input.type.tensor_type.elem_type,
[1, 28, 28, 1] # NHWC shape
)
# 替换 graph.input[0]
graph.input[0].CopyFrom(new_input)
✅ 步骤 3：插入 Transpose 节点（NHWC → NCHW）
# 创建 transpose 节点：perm=[0, 3, 1, 2] 将 (N,H,W,C) → (N,C,H,W)
transpose_node = helper.make_node(
"Transpose",
inputs=["input_nhwc"],
outputs=["input_nchw"],
perm=[0, 3, 1, 2],
name="nhwc_to_nchw"
)
# 将原图第一个节点的输入从 original_input.name 改为 "input_nchw"
first_node = graph.node[0]
for i, inp in enumerate(first_node.input):
if inp == original_input.name:
first_node.input[i] = "input_nchw"
# 将 transpose 节点插入图首部
graph.node.insert(0, transpose_node)
✅ 步骤 4：更新中间值信息（关键！避免推理报错）
ONNX 推理引擎依赖 value_info 推断维度。需显式添加 input_nchw 的类型与形状：
nchw_shape = [1, 1, 28, 28]
input_nchw_info = helper.make_tensor_value_info(
"input_nchw",
original_input.type.tensor_type.elem_type,
nchw_shape
)
graph.value_info.insert(0, input_nchw_info)
✅ 步骤 5：保存新模型并验证
# 可选：检查图结构
onnx.checker.check_model(model)
# 保存转换后的模型
onnx.save(model, "model_nhwc_input.onnx")
print("✅ Model successfully converted to accept NHWC input.")
⚠️ 注意事项与最佳实践
- 无需修改权重或算子属性：所有已有算子（如 Conv、BatchNorm）仍按 NCHW 运行，仅输入预处理被封装；
- 确保 value_info 完整：缺失 input_nchw 的 value_info 会导致某些运行时（如 ONNX Runtime）无法推断后续形状而报错；
- 批量维度灵活性：若需支持动态 batch（如 [-1, 28, 28, 1]），请将 shape 中 1 替换为 None 或 ""，并在 make_tensor_value_info 中使用 dim_param="batch_size"；
- CUDA/TensorRT 加速提示：NVIDIA 的 CUDA Execution Provider 对 NHWC 有原生优化，上述转换后配合 --use_cuda 和 --enable_memory_arena 可进一步提升性能（参考 NVIDIA ProViz 博客）；
- 验证方法：用 ONNX Runtime 加载新模型，传入 np.random.randn(1, 28, 28, 1).astype(np.float32)，确认前向无报错且输出与原模型一致。
通过以上改造，您获得了一个接口友好（NHWC）、内部标准（NCHW）、完全合规、可部署的 ONNX 模型，无需 C++ 编程或自定义算子，全程纯 Python 实现。
