Keras Tuner搜索时如何让有状态LSTM在训练轮次间重置状态

Keras Tuner没有内置的epoch状态重置参数，因为它默认封装了完整的model.fit训练流程，要实现需求可通过自定义Keras回调的方式，在每个epoch结束后主动调用reset_states()，无需修改核心调参逻辑。

实现步骤 #

• 自定义回调类，重写on_epoch_end方法
  该方法会在每个epoch的训练完成、下一轮训练开始前自动触发，直接调用当前训练模型的reset_states()即可。如果需要避免不同trial之间的状态残留，可额外重写on_train_begin方法，在单个trial的训练启动时也重置一次状态。
• 将回调传入tuner.search的fit参数中
  Keras Tuner的search方法会把额外传入的参数直接透传给内部调用的model.fit，回调列表直接作为参数传入即可。
• 修复原有代码的参数索引bug
  原代码中Dense层的循环错误复用了LSTM层循环的变量i，会导致所有Dense层绑定同一个超参数ID，需要替换为当前Dense层的循环索引j。

完整可运行代码示例 #

首先定义自定义回调：

from tensorflow.keras.callbacks import Callback

class ResetStateCallback(Callback):
    def on_train_begin(self, logs=None):
        # 每个trial训练开始前先重置一次状态
        self.model.reset_states()
    def on_epoch_end(self, epoch, logs=None):
        # 每个epoch结束后重置LSTM状态
        self.model.reset_states()

修改后的模型构建函数：

def build_model(hp):
    model = Sequential()
    # Input层的batch_shape必须和训练时的batch_size完全一致
    model.add(layers.Input(batch_shape=(batch_size, train_X.shape[1], train_X.shape[2])))
    
    # LSTM层配置
    for i in range(hp.Int('num_LSTM_layers', 1, 3)):
        model.add(layers.LSTM(
            units=hp.Int(f'units_lstm_{i}', min_value=32, max_value=512, step=4),
            batch_input_shape=(batch_size, train_X.shape[1], train_X.shape[2]),
            activation=hp.Choice(f'LSTM_activation_{i}', values=['relu','softplus','tanh', 'sigmoid','softsign','selu','elu','linear'], default='elu'),
            return_sequences=True,
            stateful=True
        ))
    
    # Dense层配置：修正原代码复用LSTM循环变量i的问题，改用当前Dense层的索引j
    for j in range(hp.Int('num_dense_layers', 1, 3)):
        model.add(layers.Dense(
            units=hp.Int(f'units_dense_{j}', min_value=64, max_value=1024, step=4),
            activation=hp.Choice(f'dense_activation_{j}', values=['relu','softplus','tanh', 'sigmoid','softsign','selu','elu','linear'], default='elu')
        ))
        model.add(layers.Dropout(rate=hp.Float(f'dropout_rate_{j}', min_value=0.01, max_value=0.50, step=0.01)))
    
    model.add(layers.Dense(train_y.shape[1], activation='linear'))
    model.compile(
        optimizer=hp.Choice('optimizers', values=['rmsprop','adam','adadelta','Nadam']),
        loss='mse',
        metrics=['mse']
    )
    return model

修改后的调参搜索调用：

tuner_bo = BayesianOptimization(
    build_model,
    objective='val_loss',
    max_trials=50,
    executions_per_trial=3,
    overwrite=True,
    num_initial_points=10,
    directory=model_path,
    project_name='LSTM_KT_2001',
    allow_new_entries=True,
    tune_new_entries=True
)

tuner_bo.search_space_summary()

# 初始化状态重置回调，传入search的参数中
reset_state_cb = ResetStateCallback()
tuner_bo.search(
    train_X, train_y,
    epochs=100,
    batch_size=batch_size, # *必须和Input层指定的batch_shape第一个维度完全一致*
    validation_data=(test_X, test_y),
    verbose=2,
    callbacks=[reset_state_cb]
)

关键注意点 #

• 必须保证tuner.search传入的batch_size和Input层指定的batch_shape的第一个维度完全一致，否则stateful LSTM会触发维度不匹配的错误。
• 该回调的执行逻辑和手动逐epoch循环训练后调用reset_states()的效果完全一致，不会干扰Keras Tuner的超参数搜索流程。
• 超参数的ID要唯一，避免不同层的参数冲突，上述代码已经通过给LSTM、Dense层的参数加前缀的方式做了规避。

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