TensorFlow/Keras自定义Transformer层模型加载报错求助

问题根源 #

当使用自定义Keras层（如TransformerBlock、TokenAndPositionEmbedding）训练并保存模型后，直接加载时Keras无法自动识别自定义层的结构和参数，导致无法正确恢复层内变量（比如报错中的dense_13变量）。

解决方法 #

1. 加载模型时指定自定义层类 #

调用load_model时，通过custom_objects参数显式传入自定义层的类定义，让Keras能够正确重建这些层：

# 确保当前脚本已定义TransformerBlock和TokenAndPositionEmbedding类，或从自定义模块导入
model = tf.keras.models.load_model(
    'FAQ_model_2.keras',
    custom_objects={
        'TransformerBlock': TransformerBlock,
        'TokenAndPositionEmbedding': TokenAndPositionEmbedding
    }
)

2. 为自定义层实现get_config方法 #

Keras保存模型时依赖层的配置信息重建实例，若自定义层未实现get_config，会导致参数丢失，进而无法恢复变量。给两个自定义层补充该方法：

TransformerBlock层补充get_config

class TransformerBlock(tf.keras.layers.Layer):
    def __init__(self, embed_dim, num_heads, ff_dim, rate=0.1, **kwargs):
        super().__init__(**kwargs)
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.ff_dim = ff_dim
        self.rate = rate
        self.att = tf.keras.layers.MultiHeadAttention(num_heads=num_heads, key_dim=embed_dim)
        self.ffn = tf.keras.Sequential(
            [tf.keras.layers.Dense(ff_dim, activation="relu"), tf.keras.layers.Dense(embed_dim),]
        )
        self.layernorm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
        self.layernorm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
        self.dropout1 = tf.keras.layers.Dropout(rate)
        self.dropout2 = tf.keras.layers.Dropout(rate)

    def get_config(self):
        config = super().get_config()
        config.update({
            'embed_dim': self.embed_dim,
            'num_heads': self.num_heads,
            'ff_dim': self.ff_dim,
            'rate': self.rate
        })
        return config

    def call(self, inputs, training=False):
        # 保留原call方法逻辑
        attn_output = self.att(inputs, inputs)
        attn_output = self.dropout1(attn_output, training=training)
        out1 = self.layernorm1(inputs + attn_output)
        ffn_output = self.ffn(out1)
        ffn_output = self.dropout2(ffn_output, training=training)
        return self.layernorm2(out1 + ffn_output)

TokenAndPositionEmbedding层补充get_config

class TokenAndPositionEmbedding(tf.keras.layers.Layer):
    def __init__(self, maxlen, vocab_size, embed_dim, **kwargs):
        super().__init__(**kwargs)
        self.maxlen = maxlen
        self.vocab_size = vocab_size
        self.embed_dim = embed_dim
        self.token_emb = tf.keras.layers.Embedding(input_dim=vocab_size, output_dim=embed_dim)
        self.pos_emb = tf.keras.layers.Embedding(input_dim=maxlen, output_dim=embed_dim)

    def get_config(self):
        config = super().get_config()
        config.update({
            'maxlen': self.maxlen,
            'vocab_size': self.vocab_size,
            'embed_dim': self.embed_dim
        })
        return config

    def call(self, x):
        # 保留原call方法逻辑
        maxlen = tf.shape(x)[-1]
        positions = tf.range(start=0, limit=maxlen, delta=1)
        positions = self.pos_emb(positions)
        x = self.token_emb(x)
        return x + positions

补充完get_config后，重新训练并保存模型，再用步骤1的方式加载，即可解决变量缺失问题。

额外排查点 #

• 确保加载模型时，自定义层的类定义和保存模型时完全一致，包括类名、参数列表、内部子层结构，不能有任何修改
• 如果使用了子类化Model类，同样需要为Model类实现get_config方法，或在加载时将Model类也传入custom_objects

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