Study.optimize 的回调函数

本教程展示了如何为 optimize() 使用和实现 Optuna 的 Callback。

Callback 在每次执行 objective 后被调用，它接收 Study 和 FrozenTrial 作为参数，并执行一些工作。

MLflowCallback 是一个很好的例子。

连续修剪（prune）一定数量的 trial 后停止优化

本示例实现了一个有状态的回调函数，如果连续修剪（prune）一定数量的 trial，则停止优化。连续修剪的 trial 数量由 threshold 指定。

import optuna


class StopWhenTrialKeepBeingPrunedCallback:
    def __init__(self, threshold: int):
        self.threshold = threshold
        self._consequtive_pruned_count = 0

    def __call__(self, study: optuna.study.Study, trial: optuna.trial.FrozenTrial) -> None:
        if trial.state == optuna.trial.TrialState.PRUNED:
            self._consequtive_pruned_count += 1
        else:
            self._consequtive_pruned_count = 0

        if self._consequtive_pruned_count >= self.threshold:
            study.stop()

这个 objective 函数会修剪除前 5 个 trial 之外的所有 trial（trial.number 从 0 开始）。

def objective(trial):
    if trial.number > 4:
        raise optuna.TrialPruned

    return trial.suggest_float("x", 0, 1)

在这里，我们将阈值设置为 2：一旦连续修剪两个 trial，优化就会结束。因此，我们预计此 study 将在 7 个 trial 后停止。

import logging
import sys

# Add stream handler of stdout to show the messages
optuna.logging.get_logger("optuna").addHandler(logging.StreamHandler(sys.stdout))

study_stop_cb = StopWhenTrialKeepBeingPrunedCallback(2)
study = optuna.create_study()
study.optimize(objective, n_trials=10, callbacks=[study_stop_cb])

A new study created in memory with name: no-name-5e5da9cd-8e79-4cbc-9e7b-60b33d9d565f
Trial 0 finished with value: 0.6218430915474218 and parameters: {'x': 0.6218430915474218}. Best is trial 0 with value: 0.6218430915474218.
Trial 1 finished with value: 0.2427726243041084 and parameters: {'x': 0.2427726243041084}. Best is trial 1 with value: 0.2427726243041084.
Trial 2 finished with value: 0.38703104273220823 and parameters: {'x': 0.38703104273220823}. Best is trial 1 with value: 0.2427726243041084.
Trial 3 finished with value: 0.6847177614644492 and parameters: {'x': 0.6847177614644492}. Best is trial 1 with value: 0.2427726243041084.
Trial 4 finished with value: 0.7591046580029978 and parameters: {'x': 0.7591046580029978}. Best is trial 1 with value: 0.2427726243041084.
Trial 5 pruned.
Trial 6 pruned.

正如你在上面的日志中看到的，study 按照预期在 7 个 trial 后停止了。