TensorFlow 2.x中因自定义梯度导致冻结EfficientNetB0保存失败

问题场景 #

在TensorFlow 2.9.2的Keras Applications中微调EfficientNetB0模型，目标是优化模型以适配TensorFlow C++ API推理。使用tensorflow.python.framework.convert_to_constants.convert_variables_to_constants_v2结合tf.saved_model.save的流程，对MobileNet、ResNet50等架构有效，但对EfficientNetBX系列模型失效。

复现代码 #

import tensorflow as tf
from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2

input_shape = (224, 224, 3)

model = tf.keras.applications.ResNet50(weights='imagenet', include_top=True, input_shape=input_shape)
# 切换为EfficientNetB0时触发错误
# model = tf.keras.applications.EfficientNetB0(weights='imagenet', include_top=True, input_shape=input_shape)

full_model = tf.function(lambda x: model(x))
full_model = full_model.get_concrete_function(tf.TensorSpec(model.inputs[0].shape,
                                                            model.inputs[0].dtype,
                                                            name='yourInputName'))
frozen_func = convert_variables_to_constants_v2(full_model)

tf.saved_model.save(frozen_func, 'saved_model')

错误信息 #

ValueError: Found invalid capture
Tensor("efficientnetb0/stem_activation/beta:0", shape=(), dtype=float32) when saving custom gradients

原因分析 #

这是TensorFlow 2.x特定版本的兼容性问题：EfficientNet的Swish激活层使用了自定义梯度实现，同时BatchNormalization层的变量（如beta）在冻结函数转换时，被自定义梯度的逻辑不当捕获，导致SavedModel保存流程出错。ResNet、MobileNet等架构未使用此类带自定义梯度的激活层，因此不受影响。

解决方法 #

方法一：先保存标准SavedModel，再转换为冻结图 #

绕过直接冻结tf.function的流程，先导出标准SavedModel，再用旧版API转换为冻结图：

import tensorflow as tf

input_shape = (224, 224, 3)
model = tf.keras.applications.EfficientNetB0(weights='imagenet', include_top=True, input_shape=input_shape)

# 保存标准SavedModel
tf.saved_model.save(model, 'temp_saved_model')

# 加载并转换为冻结图
loaded_model = tf.saved_model.load('temp_saved_model')
infer_func = loaded_model.signatures['serving_default']

# 获取图结构并冻结变量
graph_def = infer_func.graph.as_graph_def()
with tf.compat.v1.Session(graph=tf.Graph()) as sess:
    tf.import_graph_def(graph_def, name='')
    frozen_graph = tf.compat.v1.graph_util.convert_variables_to_constants(
        sess,
        sess.graph_def,
        [output.name.split(':')[0] for output in infer_func.structured_outputs.values()]
    )

# 保存冻结后的.pb文件
tf.io.write_graph(frozen_graph, './', 'frozen_efficientnetb0.pb', as_text=False)

方法二：清除自定义梯度并关闭学习阶段 #

在转换冻结函数前，清除EfficientNet激活层的自定义梯度，并强制模型进入推理模式：

import tensorflow as tf
from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2

input_shape = (224, 224, 3)
model = tf.keras.applications.EfficientNetB0(weights='imagenet', include_top=True, input_shape=input_shape)

# 强制模型进入推理模式，确保BN层使用统计值而非实时计算
tf.keras.backend.set_learning_phase(0)

# 清除激活层的自定义梯度
for layer in model.layers:
    if hasattr(layer, 'activation') and layer.activation is not None:
        if hasattr(layer.activation, '_custom_gradient'):
            delattr(layer.activation, '_custom_gradient')

full_model = tf.function(lambda x: model(x))
full_model = full_model.get_concrete_function(tf.TensorSpec(model.inputs[0].shape,
                                                            model.inputs[0].dtype,
                                                            name='yourInputName'))
frozen_func = convert_variables_to_constants_v2(full_model)

tf.saved_model.save(frozen_func, 'saved_model')

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