Optuna MedianPruner 怎样结合中间指标提前停止试验

我在训练模型时，最浪费预算的情况通常不是某一次试验报错，而是已经明显落后的参数组合仍然跑完整个训练周期。Optuna 的 MedianPruner 适合处理这种场景：objective 在每个训练步骤报告一个中间指标，剪枝器再把当前试验的进展和此前试验在同一步骤的结果比较，必要时提前结束本次试验。

官方地址：https://optuna.readthedocs.io/en/stable/reference/generated/optuna.pruners.MedianPruner.html

最短的结论是：先让指标在各个试验之间具有相同含义，再在循环中依次调用 trial.report(value, step)、trial.should_prune()，命中后抛出 optuna.TrialPruned()。如果没有稳定的中间指标，MedianPruner 没有可靠的比较依据。

一、先明确中间指标和比较方向

先把训练预算当成需要保护的资产：每个 step 都应该产生同一验证集上的同一指标，例如验证损失或验证准确率。指标必须能跨试验比较，step 也要从 0 开始递增；不要把某个试验的 batch 编号、另一试验的 epoch 编号混成一列。

如果使用验证损失，Study 应设置为 minimize，因为数值越小越好；如果使用准确率，通常使用 maximize。方向错了，剪枝器会把真正表现好的试验当成落后试验。

import optuna

# 这里把“验证损失”定义为越小越好，方向必须和指标语义一致。
study = optuna.create_study(direction="minimize")

二、配置 MedianPruner 的启动与预热

MedianPruner 的默认参数可以工作，但在训练步数较少或早期波动较大的任务里，我更愿意显式写出关键门槛。n_startup_trials 表示先完成多少个试验，达到数量前不剪枝；n_warmup_steps 表示当前试验的 step 小于该值时不剪枝，最早从这个 step 开始判断。

interval_steps 控制检查间隔，n_min_trials 则要求当前 step 至少有足够多的已报告结果参与判断。它们共同解决一个实际风险：历史样本太少、或者模型刚开始训练时指标还没有形成趋势，就急着停止。

import optuna

# 先积累 5 个完成试验；每个试验前 2 个 step 只观察，不做剪枝。
# 每隔 2 个 step 检查一次，并要求同一步骤至少有 3 个历史结果。
pruner = optuna.pruners.MedianPruner(
    n_startup_trials=5,
    n_warmup_steps=2,
    interval_steps=2,
    n_min_trials=3,
)

study = optuna.create_study(
    direction="minimize",
    pruner=pruner,
)

三、在训练循环中 report 并执行 should_prune

真正的接线点只有两个。训练或评估完成一个 step 后，用 report 保存中间指标；紧接着用 should_prune 询问剪枝器。如果返回 True，抛出 TrialPruned，Optuna 会把这个试验记录为剪枝状态，后续试验仍然可以继续。

下面的代码用一个可替换的验证函数表示模型训练，重点是调用顺序和异常类型，而不是伪造某个模型的训练结果。

import optuna

def evaluate_validation_loss(trial, step):
    # 示例用可重复的下降趋势代替真实模型评估；实际项目应返回验证集损失。
    return 1.0 / (step + 1) + trial.suggest_float("noise", 0.0, 0.02)

def objective(trial):
    # 这些参数只是示例搜索空间，训练模型时替换为真实超参数。
    learning_rate = trial.suggest_float("learning_rate", 1e-4, 1e-1, log=True)
    del learning_rate  # 保留变量表示它会被训练器使用，示例不执行真实训练。

    for step in range(12):
        validation_loss = evaluate_validation_loss(trial, step)

        # 每个 step 只上报一次可比较的验证损失，step 必须是非负整数。
        trial.report(validation_loss, step=step)

        # 命中剪枝条件就结束当前试验，不能把它当成普通异常吞掉。
        if trial.should_prune():
            raise optuna.TrialPruned()

    # 没有被剪枝时，返回最后一个完整 step 的目标值。
    return validation_loss

study = optuna.create_study(
    direction="minimize",
    pruner=optuna.pruners.MedianPruner(
        n_startup_trials=5,
        n_warmup_steps=2,
        interval_steps=2,
        n_min_trials=3,
    ),
)
study.optimize(objective, n_trials=20)

四、理解 MedianPruner 到底比较什么

它不是拿当前值和一个固定阈值比较，而是看此前已完成试验在同一步骤的中位数。对于最小化问题，如果当前试验截至当前 step 的最佳验证损失仍比历史中位数差，就可能被剪枝；对于最大化问题，比较方向相反。

这里的“同一步骤”很重要。试验 A 在 step 4 上报的值，不能直接和试验 B 的 step 7 比。若某个步骤没有足够的历史报告，n_min_trials 会阻止这次判断；如果还处于 warmup 或 startup 阶段，同样不会剪枝。

五、用间隔和最小样本数控制判断频率

如果每个 epoch 的评估很贵，可以让 interval_steps 大于 1；如果每一步都便宜且训练曲线变化快，可以缩小间隔。需要注意，检查是按 warmup 偏移后的间隔进行的，而且如果检查时尚未报告该 step，判断会推迟到有值时。

n_min_trials 适合数据波动较大的任务。它不会凭空提升模型效果，只是要求参与当前 step 比较的历史报告数量达到下限。启动试验数和最小样本数也要结合总试验数设置，不能把门槛设到几乎永远无法剪枝。

# 训练步数很多时，降低剪枝检查频率以减少额外判断。
pruner = optuna.pruners.MedianPruner(
    n_startup_trials=8,  # 先让历史基线更稳定
    n_warmup_steps=5,    # 保护每个试验最初的几个训练步骤
    interval_steps=3,    # 每 3 个 step 左右检查一次
    n_min_trials=4,      # 至少有 4 个同一步骤的历史结果
)

六、保留可审计的试验结果

最后把“提前停止”和“运行失败”分开。可以通过 Optuna 的试验状态、日志或回调记录被剪枝的试验数量，再观察完成试验的最佳值。不要捕获所有异常后统一返回一个很差的数字，否则真实训练错误会伪装成一次正常的低分试验，破坏中位数基线。

我通常会先用较保守的 n_startup_trials、n_warmup_steps 和 n_min_trials 跑一轮，再根据剪枝比例调整，而不是一开始就把剪枝器调得很激进。这样既能减少明显无望试验的计算消耗，也能保留足够的探索空间。

常见问题

MedianPruner 必须和 report 一起使用吗？

要让它基于中间进度剪枝，就必须在 objective 中调用 report，然后调用 should_prune。只配置 pruner 而不报告指标，剪枝器没有可比较的中间结果。

为什么刚开始的几个试验没有被剪枝？

这通常是 n_startup_trials 的预期行为：在指定数量的试验完成前，剪枝会被禁用。每个试验早期不剪枝则对应 n_warmup_steps。

可以把多个指标都 report 给 MedianPruner 吗？

Trial.report 面向单目标中间值。若任务是多目标优化，不要把多个目标硬塞进同一个 report；应重新设计目标或使用适配多目标的流程。