自定义ResNet模型保存失败：ExponentialMovingAverage未被追踪

一、关于“需将TensorFlow对象分配给被追踪对象的属性或主对象直接属性”的含义 #

TensorFlow保存模型时，会自动扫描并追踪所有和模型功能相关的可训练参数、资源张量、状态对象（比如EMA），但这些对象必须满足两个条件之一：

• 是自定义层/模型类的直接实例属性（比如self.ema）；
• 是已被TensorFlow追踪的对象（比如Layer的实例）的属性。

如果你的EMA对象只是在层的call方法里临时创建，或者存放在普通Python变量、未被追踪的对象中，TensorFlow的保存机制就无法感知到它的存在，当模型代码中引用了这个未被追踪的EMA对应的资源Tensor时，就会触发AssertionError。

二、针对自定义BN层EMA未被追踪的修复步骤 #

1. 将EMA对象作为自定义层的直接属性初始化 #

在自定义BatchNormalization层的__init__方法中，直接把EMA赋值给层的实例属性，让TensorFlow从初始化阶段就开始追踪它：

class CustomBatchNormalization(tf.keras.layers.Layer):
    def __init__(self, decay=0.99, **kwargs):
        super().__init__(**kwargs)
        self.decay = decay
        # 关键：将EMA设为层的直接属性
        self.ema = tf.train.ExponentialMovingAverage(decay=self.decay)

2. 绑定EMA到已追踪的BN状态变量 #

在训练模式下更新滑动均值和方差时，要让EMA作用于层的moving_mean、moving_var这些已被TensorFlow追踪的属性，确保EMA的状态和层的状态绑定：

def call(self, inputs, training=False):
    if training:
        batch_mean, batch_var = tf.nn.moments(inputs, axes=[0,1,2])
        # 让EMA追踪并更新层的moving_mean和moving_var
        self.ema.apply([self.moving_mean, self.moving_var])
        # 更新滑动统计量（需和EMA操作绑定）
        update_mean_op = tf.assign(self.moving_mean, 
                                  self.decay * self.moving_mean + (1 - self.decay) * batch_mean)
        update_var_op = tf.assign(self.moving_var, 
                                 self.decay * self.moving_var + (1 - self.decay) * batch_var)
        # 确保更新操作完成后再返回结果
        with tf.control_dependencies([update_mean_op, update_var_op]):
            return tf.nn.batch_normalization(inputs, batch_mean, batch_var, self.beta, self.gamma, 1e-5)

3. 推理时通过层属性获取EMA统计量 #

在推理模式下，必须通过层的self.ema属性来获取滑动均值和方差，不能使用未绑定的局部变量：

else:
    # 从已追踪的EMA对象中获取平均统计量
    ema_mean = self.ema.average(self.moving_mean)
    ema_var = self.ema.average(self.moving_var)
    return tf.nn.batch_normalization(inputs, ema_mean, ema_var, self.beta, self.gamma, 1e-5)

额外说明 #

如果不是必须自定义BN层，建议直接使用TensorFlow官方的tf.keras.layers.BatchNormalization，它已经内置了滑动平均逻辑，且完全兼容模型保存机制，无需手动处理EMA的追踪问题。

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