将ONNX模型转换为TensorFlow Lite时遇pytorch_half_pixel不支持错误

问题根源 #

报错明确指出coordinate_transformation_mode=pytorch_half_pixel是TensorFlow不支持的模式，这个参数通常是PyTorch导出ONNX时默认使用的，和ONNX版本无关，所以单纯更新ONNX无法解决。

解决方案：修改ONNX模型的Resize节点参数 #

直接修改ONNX模型中Resize节点的coordinate_transformation_mode属性为TensorFlow支持的类型（比如half_pixel或align_corners），步骤如下：

1. 加载并修改ONNX模型：

import onnx

# 原模型路径
onnx_model_path = '/home/sfrye/segmentation/segmentation_checkpoints/efficientnet/modified-new.onnx'
model = onnx.load(onnx_model_path)

# 遍历所有节点，定位Resize节点并修改参数
for node in model.graph.node:
    if node.op_type == "Resize":
        for attr in node.attribute:
            if attr.name == "coordinate_transformation_mode":
                # 将参数改为TensorFlow支持的"half_pixel"
                attr.s = b"half_pixel"
                print("已修改Resize节点的coordinate_transformation_mode参数")

# 保存修改后的模型
fixed_model_path = '/home/sfrye/segmentation/segmentation_checkpoints/efficientnet/modified-new-fixed.onnx'
onnx.save(model, fixed_model_path)

2. 使用修改后的模型执行转换：

import tensorflow as tf
import onnx_tf

tf_model = onnx_tf.backend.prepare(onnx.load(fixed_model_path))
tf_model.export_graph("tflite_model.tf")

# 后续转换为TFLite格式
converter = tf.lite.TFLiteConverter.from_saved_model("tflite_model.tf")
tflite_model = converter.convert()
with open("model.tflite", "wb") as f:
    f.write(tflite_model)

注意事项 #

• 如果修改为half_pixel后模型输出结果和原模型有偏差，可以尝试改为align_corners，但需要验证两种模式下的输出差异，确保符合业务需求。
• 修改模型前建议备份原ONNX文件，避免操作失误导致模型损坏。

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