TensorFlow自定义Transformer模型反序列化失败问题求助

问题重现 #

保存与加载自定义Transformer模型时执行以下代码报错：

model_directory ="saved_model.keras"
transformer_model.save(model_directory)
load_transformer_model = load_model("saved_model.keras", custom_objects=custom_objects)

核心错误信息：

TypeError: <class '__main__.Transformer'> could not be deserialized properly.
Error encountered during deserialization: Transformer.__init__() got an unexpected keyword argument 'trainable'
Received config={'trainable': True, 'dtype': 'float32'}

问题根源 #

1. 自定义Model/Layer的__init__未接收**kwargs，加载时Keras自动传入的父类参数（如trainable、dtype）无法被处理
2. 未将自定义初始化参数（如num_layers、d_model）保存为实例属性，导致get_config无法捕获这些参数
3. get_config未合并父类配置，序列化后的配置缺失实例化所需的核心参数

修复步骤 #

1. 修改Transformer类的初始化方法 #

添加**kwargs并传递给父类，同时保存所有初始化参数为实例属性：

class Transformer(tf.keras.Model):
    """
    The Transformer model architecture, consisting of an Encoder and Decoder.
    """

    def __init__(
        self,
        num_layers,
        d_model,
        num_heads,
        d_feedforward,
        input_vocab_size,
        target_vocab_size,
        max_num_positions_in_pe_encoder,
        max_num_positions_in_pe_decoder,
        dropout_rate=0.1,
        **kwargs  # 新增：接收父类传递的参数
     ):
        super(Transformer, self).__init__(**kwargs)  # 传递参数给父类
        # 保存所有初始化参数为实例属性，供get_config调用
        self.num_layers = num_layers
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_feedforward = d_feedforward
        self.input_vocab_size = input_vocab_size
        self.target_vocab_size = target_vocab_size
        self.max_num_positions_in_pe_encoder = max_num_positions_in_pe_encoder
        self.max_num_positions_in_pe_decoder = max_num_positions_in_pe_decoder
        self.dropout_rate = dropout_rate

        self.encoder = Encoder(
            num_layers=num_layers,
            d_model=d_model,
            num_heads=num_heads,
            d_feedforward=d_feedforward,
            input_vocab_size=input_vocab_size,
            maximum_positions_in_pe=max_num_positions_in_pe_encoder,
            dropout_rate=dropout_rate,
            **kwargs  # 传递参数给Encoder层
        )
        self.decoder = Decoder(
            num_layers=num_layers,
            d_model=d_model,
            num_heads=num_heads,
            d_feedforward=d_feedforward,
            target_vocab_size=target_vocab_size,
            maximum_positions_in_pe=max_num_positions_in_pe_decoder,
            dropout_rate=dropout_rate,
            **kwargs  # 传递参数给Decoder层
        )
        self.final_layer = Dense(target_vocab_size)

2. 重写get_config方法 #

合并父类配置与自定义参数，确保序列化的配置包含所有实例化必要信息：

def get_config(self):
    base_config = super(Transformer, self).get_config()
    custom_config = {
        'num_layers': self.num_layers,
        'd_model': self.d_model,
        'num_heads': self.num_heads,
        'd_feedforward': self.d_feedforward,
        'input_vocab_size': self.input_vocab_size,
        'target_vocab_size': self.target_vocab_size,
        'max_num_positions_in_pe_encoder': self.max_num_positions_in_pe_encoder,
        'max_num_positions_in_pe_decoder': self.max_num_positions_in_pe_decoder,
        'dropout_rate': self.dropout_rate,
    }
    return {**base_config, **custom_config}

3. 重写from_config方法（可选） #

若默认实例化逻辑不满足需求，可自定义实例化规则：

@classmethod
def from_config(cls, config):
    return cls(**config)

4. 同步修改所有自定义层 #

Encoder、Decoder、EncoderLayer、DecoderLayer需执行相同修改：

• 在__init__中添加**kwargs并传递给父类
• 保存所有初始化参数为实例属性
• 重写get_config方法合并父类与自定义配置

5. 保存与加载模型 #

确保custom_objects包含所有自定义类：

custom_objects = {
    'Transformer': Transformer,
    'Encoder': Encoder,
    'Decoder': Decoder,
    'EncoderLayer': EncoderLayer,
    'DecoderLayer': DecoderLayer
}

model_directory ="saved_model.keras"
transformer_model.save(model_directory)
load_transformer_model = load_model(model_directory, custom_objects=custom_objects)

验证方法 #

1. 重新初始化或训练模型
2. 保存模型后加载，测试预测功能，确认无报错且输出结果与原模型一致

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