使用TensorFlow子类化API实现Unet报OperatorNotAllowedInGraphError如何解决

报错根因 #

你遇到的OperatorNotAllowedInGraphError是TensorFlow图执行模式的典型限制：图模式下不支持直接遍历tf.Tensor对象，只有即时执行（Eager Execution）模式默认支持该操作。该错误通常是call()方法中使用了Python原生的非图兼容语法导致的。

常见触发点排查 #

你可以对照自己的实现检查以下场景：

• 直接用Python的for循环迭代张量，比如遍历多通道特征、遍历不同尺度的特征张量，这类操作要替换为TensorFlow内置算子，比如拼接跳层特征直接用tf.keras.layers.Concatenate(axis=-1)([low_level_feature, upsampled_feature])，不要循环遍历每个特征再处理
• 用Python原生的if判断直接以张量值作为判断条件，这类场景要替换为tf.cond实现分支逻辑
• 用len(tensor)、tensor[i]等Python原生语法操作张量维度/索引，获取张量维度要用tf.shape(tensor)[dim_idx]，切片操作要用tf.slice实现
• 如果你的下采样/上采样模块存放在Python列表中，需要确保该列表是在__init__方法中定义为类的成员属性（比如self.down_blocks = [DownBlock(64), DownBlock(128), ...]），在call中遍历类成员属性的层列表是合法的，不要在call内部临时生成层列表再遍历

临时验证方法 #

可以在代码入口处添加一行tf.config.run_functions_eagerly(True)，强制全局开启即时执行模式。如果开启后代码可以正常运行、输出结果符合预期，即可确认是图兼容语法问题，逐行替换为TensorFlow内置算子即可。注意不要长期保留该配置，会大幅降低训练效率。

call()方法逻辑校验 #

标准Unet的call方法逻辑应该符合以下流程，你可以对照修正：

1. 输入张量依次传入所有下采样块，每一个下采样块的输出都单独存储，作为后续跳层连接的特征
2. 最后一个下采样块的输出传入瓶颈层处理
3. 瓶颈层输出依次传入所有上采样块，每个上采样块的输入为「上一层上采样的输出」和「对应尺度的下采样跳层特征」拼接后的结果
4. 最后一个上采样块的输出传入输出卷积层，得到和输入尺寸匹配的分割预测结果

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