基于Ray Tune PBT实现条件搜索空间：优化模型隐藏层数

问题描述 #

我想要为以下模型找到最佳隐藏层数，当前模型使用Conv1D和LSTM的堆叠结构，同时用Ray Tune结合Population Based Training（PBT）进行超参数调优。现在的问题是：当Conv1DAn/Conv1DBn/LSTMn设置为0或某个数值时，对应的后续层参数（比如Conv1DAn_1、LSTMn_2等）其实是未被使用的，我希望实现条件搜索空间，只在需要的时候才搜索这些未被使用的层参数，避免无效的参数采样。

当前模型代码 #

x1 = Input((window_size, 3), name='x1')
x2 = Input((window_size, 3), name='x2')
convA1 = Conv1D(config["Conv1DA"],11,padding='same',activation='relu')(x1)
for i in range(config["Conv1DAn"]):
    if i > 0: 
        convA2 = Conv1D(config[f'Conv1DAn_{i}'],11,padding='same',activation='relu')(convA1)
poolA = MaxPooling1D(3)(convA1)

convB1 = Conv1D(config["Conv1DB"],11,padding='same',activation='relu')(x2)
for i in range(config["Conv1DBn"]):
    if i > 0:
        convB2 = Conv1D(config[f'Conv1DBn_{i}'],11,padding='same',activation='relu')(convB1)
poolB = MaxPooling1D(3)(convB1)
AB = concatenate([poolA, poolB])

lstm1 = Bidirectional(LSTM(config["LSTM1"], return_sequences=True))(AB)
drop1 = Dropout(config['dropout'])(lstm1)
for i in range(config['LSTMn']):
    if i > 0:
        lstm2 = Bidirectional(LSTM(config[f'LSTMn_{i}'], return_sequences=True))(drop1)
        drop1 =  Dropout(config['dropout'])(lstm2)   
lstm2 = Bidirectional(LSTM(config['LSTMn_l']))(drop1)
drop2 =  Dropout(config['dropout'])(lstm2)
y1_pred = Dense(4,kernel_regularizer='l2')(drop2)
model = Model(inputs =[x1, x2], outputs = [y1_pred])

当前Ray Tune + PBT调参代码 #

if __name__ == "__main__":
    pbt = PopulationBasedTraining(
        perturbation_interval=600,
        hyperparam_mutations={
            "dropout": tune.uniform(0.1,0.5),
            "lr": tune.uniform(1e-5,1e-3),
            "Conv1DA": tune.randint(10,50),
            "Conv1DAn": tune.choice([0,1,2,3]),
            "Conv1DAn_1": tune.randint(10,50),
            "Conv1DAn_2": tune.randint(10,50),
            "Conv1DAn_3": tune.randint(10,50),
            "Conv1DB": tune.randint(10,50),
            "Conv1DBn": tune.choice([0,1,2,3]),
            "Conv1DBn_1": tune.randint(10,50),
            "Conv1DBn_2": tune.randint(10,50),
            "Conv1DBn_3": tune.randint(10,50),
            "LSTM1": tune.randint(10,50),
            "LSTMn": tune.choice([0,1,2,3]),
            "LSTMn_1": tune.randint(10,50),
            "LSTMn_2": tune.randint(10,50),
            "LSTMn_3": tune.randint(10,50),
            "LSTMn_l": tune.randint(10,50)})
    tuner = tune.Tuner(
        BroadModel,
        run_config=air.RunConfig(
            name="BroadPBT"+timestr,
            stop={"training_iteration": 4}),
        tune_config=tune.TuneConfig(
            reuse_actors=True,
            scheduler=pbt,
            metric="loss",
            mode="min",
            num_samples=2),
        param_space={
            "finish_fast": False,
            "batch_size": 1000,
            "epochs": 4       ,
            "dropout": tune.uniform(0.1,0.5),
            "lr": tune.uniform(1e-5,1e-3),
            "Conv1DA": tune.randint(10,50),
            "Conv1DAn": tune.choice([0,1,2,3]),
            "Conv1DAn_1": tune.randint(10,50),
            "Conv1DAn_2": tune.randint(10,50),
            "Conv1DAn_3": tune.randint(10,50),
            "Conv1DB": tune.randint(10,50),
            "Conv1DBn": tune.choice([0,1,2,3]),
            "Conv1DBn_1": tune.randint(10,50),
            "Conv1DBn_2": tune.randint(10,50),
            "Conv1DBn_3": tune.randint(10,50),
            "LSTM1": tune.randint(10,50),
            "LSTMn": tune.choice([0,1,2,3]),
            "LSTMn_1": tune.randint(10,50),
            "LSTMn_2": tune.randint(10,50),
            "LSTMn_3": tune.randint(10,50),
            "LSTMn_l": tune.randint(10,50)})

解决方案 #

要实现条件搜索空间，需要同时修改模型构建逻辑和Ray Tune的参数空间配置，确保只有当层数参数（如Conv1DAn）大于0时，才会生成对应的层并搜索其参数。

1. 修复模型构建逻辑 #

原模型的循环逻辑存在层数不匹配问题，比如Conv1DAn=2时实际只生成1层，先修正逻辑让层数参数和实际堆叠层数一致：

x1 = Input((window_size, 3), name='x1')
x2 = Input((window_size, 3), name='x2')

# 处理ConvA分支
convA = Conv1D(config["Conv1DA"], 11, padding='same', activation='relu')(x1)
# 根据Conv1DAn的值堆叠对应数量的Conv1D层
for i in range(1, config["Conv1DAn"] + 1):
    convA = Conv1D(config[f'Conv1DAn_{i}'], 11, padding='same', activation='relu')(convA)
poolA = MaxPooling1D(3)(convA)

# 处理ConvB分支
convB = Conv1D(config["Conv1DB"], 11, padding='same', activation='relu')(x2)
for i in range(1, config["Conv1DBn"] + 1):
    convB = Conv1D(config[f'Conv1DBn_{i}'], 11, padding='same', activation='relu')(convB)
poolB = MaxPooling1D(3)(convB)

AB = concatenate([poolA, poolB])

# 处理LSTM分支
lstm_out = Bidirectional(LSTM(config["LSTM1"], return_sequences=True))(AB)
lstm_out = Dropout(config['dropout'])(lstm_out)
# 根据LSTMn的值堆叠对应数量的LSTM层
for i in range(1, config["LSTMn"] + 1):
    lstm_out = Bidirectional(LSTM(config[f'LSTMn_{i}'], return_sequences=True))(lstm_out)
    lstm_out = Dropout(config['dropout'])(lstm_out)
# 最后一层LSTM
lstm_out = Bidirectional(LSTM(config['LSTMn_l']))(lstm_out)
lstm_out = Dropout(config['dropout'])(lstm_out)

y1_pred = Dense(4, kernel_regularizer='l2')(lstm_out)
model = Model(inputs=[x1, x2], outputs=[y1_pred])

2. 配置Ray Tune的条件搜索空间 #

使用Ray Tune的tune.with_condition函数定义条件参数，只有当对应的层数参数满足条件时，才会启用该参数的搜索，同时同步修改PBT的超参数突变配置：

from ray import tune
from ray.tune.schedulers import PopulationBasedTraining
from ray import air

if __name__ == "__main__":
    # 定义条件参数空间
    param_space = {
        "finish_fast": False,
        "batch_size": 1000,
        "epochs": 4,
        "dropout": tune.uniform(0.1, 0.5),
        "lr": tune.uniform(1e-5, 1e-3),
        # ConvA相关参数
        "Conv1DA": tune.randint(10, 50),
        "Conv1DAn": tune.choice([0, 1, 2, 3]),
        # 只有当Conv1DAn >=1时，才搜索Conv1DAn_1
        "Conv1DAn_1": tune.with_condition(
            tune.randint(10, 50),
            lambda config: config["Conv1DAn"] >= 1
        ),
        # 只有当Conv1DAn >=2时，才搜索Conv1DAn_2
        "Conv1DAn_2": tune.with_condition(
            tune.randint(10, 50),
            lambda config: config["Conv1DAn"] >= 2
        ),
        # 只有当Conv1DAn >=3时，才搜索Conv1DAn_3
        "Conv1DAn_3": tune.with_condition(
            tune.randint(10, 50),
            lambda config: config["Conv1DAn"] >= 3
        ),
        # ConvB相关参数
        "Conv1DB": tune.randint(10, 50),
        "Conv1DBn": tune.choice([0, 1, 2, 3]),
        "Conv1DBn_1": tune.with_condition(
            tune.randint(10, 50),
            lambda config: config["Conv1DBn"] >= 1
        ),
        "Conv1DBn_2": tune.with_condition(
            tune.randint(10, 50),
            lambda config: config["Conv1DBn"] >= 2
        ),
        "Conv1DBn_3": tune.with_condition(
            tune.randint(10, 50),
            lambda config: config["Conv1DBn"] >= 3
        ),
        # LSTM相关参数
        "LSTM1": tune.randint(10, 50),
        "LSTMn": tune.choice([0, 1, 2, 3]),
        "LSTMn_1": tune.with_condition(
            tune.randint(10, 50),
            lambda config: config["LSTMn"] >= 1
        ),
        "LSTMn_2": tune.with_condition(
            tune.randint(10, 50),
            lambda config: config["LSTMn"] >= 2
        ),
        "LSTMn_3": tune.with_condition(
            tune.randint(10, 50),
            lambda config: config["LSTMn"] >= 3
        ),
        "LSTMn_l": tune.randint(10, 50)
    }

    # 定义PBT的条件突变配置
    pbt = PopulationBasedTraining(
        perturbation_interval=600,
        hyperparam_mutations={
            "dropout": tune.uniform(0.1, 0.5),
            "lr": tune.uniform(1e-5, 1e-3),
            "Conv1DA": tune.randint(10, 50),
            "Conv1DAn": tune.choice([0, 1, 2, 3]),
            "Conv1DAn_1": tune.with_condition(
                tune.randint(10, 50),
                lambda config: config["Conv1DAn"] >= 1
            ),
            "Conv1DAn_2": tune.with_condition(
                tune.randint(10, 50),
                lambda config: config["Conv1DAn"] >= 2
            ),
            "Conv1DAn_3": tune.with_condition(
                tune.randint(10, 50),
                lambda config: config["Conv1DAn"] >= 3
            ),
            "Conv1DB": tune.randint(10, 50),
            "Conv1DBn": tune.choice([0, 1, 2, 3]),
            "Conv1DBn_1": tune.with_condition(
                tune.randint(10, 50),
                lambda config: config["Conv1DBn"] >= 1
            ),
            "Conv1DBn_2": tune.with_condition(
                tune.randint(10, 50),
                lambda config: config["Conv1DBn"] >= 2
            ),
            "Conv1DBn_3": tune.with_condition(
                tune.randint(10, 50),
                lambda config: config["Conv1DBn"] >= 3
            ),
            "LSTM1": tune.randint(10, 50),
            "LSTMn": tune.choice([0, 1, 2, 3]),
            "LSTMn_1": tune.with_condition(
                tune.randint(10, 50),
                lambda config: config["LSTMn"] >= 1
            ),
            "LSTMn_2": tune.with_condition(
                tune.randint(10, 50),
                lambda config: config["LSTMn"] >= 2
            ),
            "LSTMn_3": tune.with_condition(
                tune.randint(10, 50),
                lambda config: config["LSTMn"] >= 3
            ),
            "LSTMn_l": tune.randint(10, 50)
        }
    )

    tuner = tune.Tuner(
        BroadModel,
        run_config=air.RunConfig(
            name="BroadPBT"+timestr,
            stop={"training_iteration": 4}
        ),
        tune_config=tune.TuneConfig(
            reuse_actors=True,
            scheduler=pbt,
            metric="loss",
            mode="min",
            num_samples=2
        ),
        param_space=param_space
    )
    # 启动调优
    results = tuner.fit()

关键说明 #

• tune.with_condition会通过lambda表达式判断是否启用参数：比如Conv1DAn=0时，Conv1DAn_1/Conv1DAn_2/Conv1DAn_3不会被纳入搜索空间，避免无效采样。
• 修正后的模型循环逻辑确保Conv1DAn=N时，会准确堆叠N个额外的Conv1D层，和参数定义完全匹配。
• PBT的超参数突变配置同步加入条件判断，避免突变时生成无效的参数组合。

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