将ONNX模型输入布局从NCHW转换为NHWC的方法

完全可以修改这类ONNX模型，让它接受NHWC布局的输入。核心思路是在模型输入节点后插入Transpose算子，将NHWC格式的输入转换为模型原本需要的NCHW格式，同时调整输入的形状定义。以下是具体实现步骤：

步骤1：加载模型并确认输入信息 #

先加载模型，查看原输入的名称和形状：

import onnx

model_path = "your_model.onnx"
model_obj = onnx.load(model_path)

# 获取原输入节点
input_tensor = model_obj.graph.input[0]
print(f"原输入名称: {input_tensor.name}")
print(f"原输入形状: {[dim.dim_value for dim in input_tensor.type.tensor_type.shape.dim]}")

步骤2：修改输入的形状定义 #

将原输入形状从(1,1,28,28)调整为(1,28,28,1)：

# 重置输入形状参数
input_shape = input_tensor.type.tensor_type.shape.dim
input_shape[0].dim_value = 1
input_shape[1].dim_value = 28
input_shape[2].dim_value = 28
input_shape[3].dim_value = 1

步骤3：插入Transpose算子转换布局 #

在输入节点和原模型首个计算节点之间插入转置算子，完成NHWC（维度顺序0,1,2,3）到NCHW（维度顺序0,3,1,2）的转换：

from onnx import helper

# 创建Transpose节点
transpose_node = helper.make_node(
    "Transpose",
    inputs=[input_tensor.name],
    outputs=["transposed_input"],
    perm=[0, 3, 1, 2]  # 对应NHWC转NCHW的维度置换顺序
)

# 将转置节点插入到模型计算节点的最前端
model_obj.graph.node.insert(0, transpose_node)

# 替换原模型中所有使用原输入名称的节点，改为转置后的输出名称
for node in model_obj.graph.node[1:]:
    for i, input_name in enumerate(node.input):
        if input_name == input_tensor.name:
            node.input[i] = "transposed_input"

步骤4：验证并保存修改后的模型 #

完成修改后，验证模型结构合法性，再保存新模型：

# 检查模型结构是否符合ONNX规范
onnx.checker.check_model(model_obj)

# 保存修改后的模型
output_model_path = "model_nhwc.onnx"
onnx.save(model_obj, output_model_path)

注意事项 #

• 若模型存在多个输入，需针对每个需要转换布局的输入重复上述操作。
• 转置算子不会引入额外计算开销，多数推理框架会自动对这类算子做优化。
• 如果原模型输入包含动态维度（如batch_size设为-1），修改形状时要对应调整动态维度的位置。

内容的提问来源于stack exchange，提问作者moriaz