tf.function单例变量错误：自定义Keras自编码器训练报错求助

核心问题 #

错误根源是在call方法内动态创建带可训练变量的层（比如BatchNormalization、SpatialDropout2D）。tf.function（Keras的fit会自动用它包装训练逻辑）要求所有tf.Variable只能创建一次，而每次调用call时重新实例化层，会重复生成变量，触发报错。

修复步骤 #

把所有带变量/需要复用的层移到类的__init__方法中初始化，确保变量只创建一次，在call方法中仅调用已初始化的层实例。

修改后的完整代码 #

class Encoder(tf.keras.Model):
    def __init__(self, params):
        super(Encoder, self).__init__()
        self.params = params
        
        # 初始化所有层（包括BN、Dropout）
        self.hidden_layer_one = tf.keras.layers.Conv2D(
            filters=self.params['e_filters'],
            kernel_size=self.params['e_size'],
            activation=self.params['activation'], 
            strides=1, 
            padding='same',
            kernel_initializer=params['initializer'], 
            use_bias=False
        )
        self.bn1 = tf.keras.layers.BatchNormalization()
        self.dropout1 = tf.keras.layers.SpatialDropout2D(0.2)
        
        self.hidden_layer_two = tf.keras.layers.Conv2D(
            filters=self.params['num_endmembers'], 
            kernel_size=1,
            activation=self.params['activation'], 
            strides=1, 
            padding='same',
            kernel_initializer=self.params['initializer'], 
            use_bias=False
        )
        self.bn2 = tf.keras.layers.BatchNormalization()
        self.dropout2 = tf.keras.layers.SpatialDropout2D(0.2)
        
        self.asc_layer = ASC(params=self.params, name='abundances')

    def call(self, input_patch):
        code = self.hidden_layer_one(input_patch)
        code = self.bn1(code)
        code = self.dropout1(code)
        code = self.hidden_layer_two(code)
        code = self.bn2(code)
        code = self.dropout2(code)
        code = self.asc_layer(code)
        return code

class Autoencoder(tf.keras.Model):
    def __init__(self, params):
        super(Autoencoder, self).__init__()
        self.encoder = Encoder(params)
        self.decoder = Decoder(params)
        self.params = params
        # 初始化Dropout层
        self.dropout_enc = tf.keras.layers.SpatialDropout2D(0.08)

    def call(self, patch):
        abunds = self.encoder(patch)
        abunds = self.dropout_enc(abunds)
        output = self.decoder(abunds)
        return output

    def train(self, patches, callback):
        self.plotWhileTraining = callback
        self.fit(patches, patches, epochs=self.params['epochs'], batch_size=self.params['batch_size'],
                 callbacks=[self.plotWhileTraining], verbose=0)

关键说明 #

1. BatchNormalization包含可训练变量（gamma、beta），必须在__init__中初始化，不能在call里动态创建
2. SpatialDropout2D虽无训练变量，但作为层实例，也应提前初始化保证复用性，避免每次调用生成新的层对象
3. 无需手动调用build()：当第一次传入训练数据时，Keras会自动根据输入形状构建所有层的权重

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