使用tf2onnx 1.12.1转TensorFlow模型至ONNX遇ExtractImagePatches不支持错误

问题场景 #

使用tf2onnx 1.12.1转换包含BatchNormalization层的TensorFlow模型到ONNX格式时触发错误，问题出在模型初始层的自定义Involution操作环节。

模型层实现代码 #

# Layer 1
x, _ = Involution(channel=3, group_number=1, kernel_size=3, stride=1, reduction_ratio=2, name="inv_rgb")(input_image, dep_image)

x = BatchNormalization(name='norm_1')(x)
x = LeakyReLU(alpha=0.1)(x)
x = MaxPooling2D(pool_size=(2, 2))(x)


y, _ = Involution_depth(channel=3, group_number=1, kernel_size=3, stride=2, reduction_ratio=1, name="inv_dep")(dep_image)
#y = Conv2D(32, (3,3), strides=(1,1), padding='same', name='dconv_1', use_bias=False)(dep_image)

y = BatchNormalization(name='dnorm_1')(y)
y = LeakyReLU(alpha=0.1)(y)
y = MaxPooling2D(pool_size=(2, 2))(y)

模型转换代码 #

import onnx
import tf2onnx

# 导出模型为ONNX格式
model.load_weights('./weightModel.h5')
tf.saved_model.save(model, "tmp_model")

!python -m tf2onnx.convert --saved-model tmp_model --output "weightFinal12.onnx"

报错信息 #

错误 - TensorFlow算子 [StatefulPartitionedCall/functional_1/inv_rgb/ExtractImagePatches: ExtractImagePatches] 不被支持

解决方法 #

• 升级tf2onnx版本：tf2onnx 1.12.1属于旧版本，后续版本已支持ExtractImagePatches算子，执行以下命令升级：

  pip install --upgrade tf2onnx
  

• 自定义算子转换规则：若升级后仍有问题，需为自定义Involution层内的ExtractImagePatches编写专属转换逻辑，参考tf2onnx自定义算子转换规范实现Converter类，转换时通过--custom-ops参数指定脚本路径
• 替换等效操作：如果自定义转换成本较高，可尝试用ONNX原生支持的算子组合（如切片+拼接）模拟ExtractImagePatches的图像分块提取功能

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