ONNX转TFLite过程中Add节点维度不匹配问题求助

1. 修正PyTorch转ONNX时的维度对齐 #

• 导出ONNX时强制统一为TensorFlow/TFLite默认的NHWC格式，避免PyTorch原生NCHW格式带来的维度混乱，同时指定兼容的opset版本：

  import torch
  model = torch.load("your_model.pth")
  model.eval()
  # 输入形状按NHWC顺序定义
  dummy_input = torch.randn(2, 384, 768, 64)
  torch.onnx.export(
      model,
      dummy_input,
      "model.onnx",
      opset_version=17,
      input_names=["input"],
      output_names=["output"],
      dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}}
  )
  

• 导出后用onnx.checker.check_model("model.onnx")验证模型合法性，再用Netron可视化工具查看Add节点的输入张量维度，确认导出阶段是否已存在维度翻转问题。

2. 优化onnx2tf的参数配置 #

• 开启全局NHWC转换，并针对目标Add节点单独配置维度修正，明确指定需要转置的输入索引和维度顺序：

  from onnx2tf import convert
  
  convert(
      input_onnx_file_path="model.onnx",
      output_folder_path="tflite_output",
      convert_to_nhwc=True,  # 全局强制NHWC转换
      param_replacement=[
          {
              "op_name": "/fnet/layer1/layer1.1/Add",
              "input_index": 0,  # 根据实际情况调整0或1
              "transpose": [0, 2, 1, 3]  # 将[2,768,384,64]转为[2,384,768,64]
          }
      ]
  )
  

• 注意：需确认Add节点的两个输入中哪个维度异常，对应调整input_index和transpose的维度参数，确保两者形状完全匹配。

3. 调整版本兼容性 #

• 尝试将onnx-tf版本降至1.16.0（与当前onnx 1.16.1版本匹配），高版本onnx-tf可能对部分ONNX算子处理存在兼容性问题：

  pip uninstall onnx-tf -y
  pip install onnx-tf==1.16.0
  

• 若仍有问题，可临时切换到TensorFlow 2.15.0稳定版测试，避免新版本的适配bug。

4. 手动修正ONNX模型 #

• 直接修改ONNX模型，在Add节点前插入Transpose算子修正维度：

  import onnx
  from onnx import helper
  
  model = onnx.load("model.onnx")
  graph = model.graph
  
  # 定位目标Add节点
  add_node = next(node for node in graph.node if node.name == "/fnet/layer1/layer1.1/Add")
  
  # 创建Transpose节点修正维度
  transpose_node = helper.make_node(
      "Transpose",
      inputs=[add_node.input[0]],
      outputs=["transposed_input"],
      perm=[0, 2, 1, 3]
  )
  # 将Transpose节点插入到Add节点之前
  graph.node.insert(graph.node.index(add_node), transpose_node)
  
  # 更新Add节点的输入为转置后的张量
  add_node.input[0] = "transposed_input"
  
  # 保存修正后的模型
  onnx.save(model, "fixed_model.onnx")
  

• 用修正后的模型再执行onnx2tf转换。

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