自定义Attention层深度学习模型加载后性能骤降的解决求助

我训练了一个带有自定义Attention层的深度学习模型，训练完成后在测试集上的准确率约为90.4%，但关闭Colab重新加载模型后，同一测试集上的性能变得极差。

自定义Attention层代码 #

class Attention(Layer):
    def __init__(self, **kwargs):
        super(Attention, self).__init__(**kwargs)

    def build(self, input_shape):
        # Initialize weights for attention
        self.W = self.add_weight(name='attention_weights', shape=(input_shape[-1], 1), initializer='uniform', trainable=True)
        super(Attention, self).build(input_shape)

    def call(self, inputs):
        # Compute attention scores and weights
        e = K.tanh(K.dot(inputs, self.W))
        a = K.softmax(e, axis=1)
        weighted = inputs * a
        # Compute weighted average of inputs
        attention = K.sum(weighted, axis=1)
        return attention
        
    def get_config(self):
        config = super(Attention, self).get_config()
        return config

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

模型结构代码 #

inputs = Input(shape=(x_train.shape[1:]))
x = Conv1D(filters=128, kernel_size=3, activation='relu')(inputs)
x = Conv1D(filters=128, kernel_size=3, activation='relu')(x)
x = MaxPooling1D(pool_size=3)(x)
x = Dropout(0.2)(x)

x = Bidirectional(GRU(units=128, activation='tanh', return_sequences=True))(x)
x = Dropout(0.2)(x)
x = Conv1D(filters=128, kernel_size=3, activation='relu')(x)
x = MaxPooling1D(pool_size=3)(x)
x = Dropout(0.2)(x)

x = Bidirectional(GRU(units=128, activation='tanh', return_sequences=True))(x)
x = Dropout(0.2)(x)
x = Conv1D(filters=128, kernel_size=3, activation='relu')(x)
x = MaxPooling1D(pool_size=3)(x)
x = Dropout(0.2)(x)

x = Bidirectional(GRU(units=128, activation='tanh', return_sequences=True))(x)
x = Attention()(x)  # Using the custom Attention layer here
x = Flatten()(x)
x = Dense(units=64, activation='relu')(x)
x = Dropout(0.2)(x)
x = Dense(units=64, activation='relu')(x)
outputs = Dense(units=2, activation='softmax')(x)

model = Model(inputs=inputs, outputs=outputs)

模型编译、训练与保存代码 #

编译代码 #

model.compile(optimizer="adam", loss="binary_crossentropy", metrics=['accuracy'])

训练与保存代码 #

checkpoint = ModelCheckpoint("modelliGRU/model_GRU_new_V3_bestWEIGHTS_FINAL_binary.h5", 
                             monitor="val_accuracy", # Metric to monitor
                             save_best_only=True, # Save only the best model
                             save_weights_only=False, # Save the entire model
                             mode='max', 
                             verbose=1)

def lr_schedule(epoch, lr):
    if epoch > 70 and \
            (epoch - 1) % 10 == 0:
        lr *= 0.1
    print("Learning rate: ", lr)
    return lr

lr_scheduler = LearningRateScheduler(lr_schedule) # Dynamic adjustment learning rate

history = model.fit(x_train, y_train, batch_size=128, epochs=100, validation_data=(x_val,y_val), callbacks=(checkpoint,lr_scheduler))
model.save(os.path.join("modelliGRU/model_GRU_new_v3_FINAL_binary.h5")) # Save training model

模型加载代码 #

model = tf.keras.models.load_model('modelliGRU/model_GRU_new_V3_bestWEIGHTS_FINAL_binary.h5',custom_objects={'Attention':Attention})

可能的原因及解决方法 #

• 自定义层权重未正确加载：虽然实现了get_config和from_config，但加载时可能存在权重恢复异常。可以：
  1. 加载模型后，打印Attention层的权重形状和值，与训练保存前的结果对比，确认W权重是否正确恢复：

     print(model.get_layer('attention').W.numpy())
     

  2. 尝试先构建完整模型结构，再单独加载权重：

     # 重新定义模型结构（需和训练时完全一致）
     inputs = Input(shape=(x_train.shape[1:]))
     # ... 完整模型结构代码 ...
     model = Model(inputs=inputs, outputs=outputs)
     # 仅加载权重
     model.load_weights('modelliGRU/model_GRU_new_V3_bestWEIGHTS_FINAL_binary.h5')
     

• 数据预处理不一致：重新打开Colab后，测试集的预处理流程可能和训练时不同，比如归一化参数未复用、数据维度或类型不匹配。需确保：
  1. 复现训练时的所有预处理步骤，比如使用训练时保存的均值、标准差进行归一化，而非重新计算测试集的统计量。
  2. 确认测试集的形状、数据类型（如float32）与训练集完全一致。
• Dropout层行为异常：训练时Dropout启用，测试时Keras默认会自动切换到推理模式，但可以手动确认：

  # 手动设置为推理模式
  model.trainable = False
  # 或使用evaluate时明确设置
  model.evaluate(x_test, y_test, batch_size=128)
  

• 模型保存格式兼容性问题：H5格式在Colab中可能存在兼容性问题，改用SavedModel格式保存和加载：

  # 保存时
  model.save("modelliGRU/model_GRU_new_v3_FINAL_binary")  # 无需.h5后缀
  # 加载时
  model = tf.keras.models.load_model("modelliGRU/model_GRU_new_v3_FINAL_binary", custom_objects={'Attention':Attention})
  

• 自定义层axis参数错误：检查softmax的axis参数是否匹配输入维度。代码中axis=1是对时间步维度做softmax，需确认测试集输入的维度顺序与训练时一致（如(batch_size, timesteps, features)），避免维度错位导致注意力计算异常。

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