用户定义的 Pruner

在 optuna.pruners 中, 我们描述了一个目标函数如何可选地包含一些 pruning 函数调用, 这些函数允许 Optuna 去终结中间结果无望的 tiral. 在本文档中, 我们描述了如何实现一个你自己的 pruner, 也即一个自定义 trial 终止条件的策略.

Pruning 界面概述

The create_study() constructor takes, as an optional argument, a pruner inheriting from BasePruner. The pruner should implement the abstract method prune(), which takes arguments for the associated Study and Trial and returns a boolean value: True if the trial should be pruned and False otherwise. Using the Study and Trial objects, you can access all other trials through the get_trial() method and, and from a trial, its reported intermediate values through the intermediate_values() (a dictionary which maps an integer step to a float value).

你可以参考 Optuna 内置的 pruner 的源代码作为构建你自己的 pruner 的模板. 在本文档中, 我们描述了一个简单 (但是却激进) 的 pruner 实现过程, 它会对那些处于同步骤 trial 中最后一名的 trial 进行剪枝.

备注

Please refer to the documentation of BasePruner or, for example, ThresholdPruner or PercentilePruner for more robust examples of pruner implementation, including error checking and complex pruner-internal logic.

例子: 实现 LastPlacePruner

We aim to optimize the loss and alpha hyperparameters for a stochastic gradient descent classifier (SGDClassifier) run on the sklearn iris dataset. We implement a pruner which terminates a trial at a certain step if it is in last place compared to completed trials at the same step. We begin considering pruning after a “warmup” of 1 training step and 5 completed trials. For demonstration purposes, we print() a diagnostic message from prune when it is about to return True (indicating pruning).

It may be important to note that the SGDClassifier score, as it is evaluated on a holdout set, decreases with enough training steps due to overfitting. This means that a trial could be pruned even if it had a favorable (high) value on a previous training set. After pruning, Optuna will take the intermediate value last reported as the value of the trial.

import numpy as np
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.linear_model import SGDClassifier

import optuna
from optuna.pruners import BasePruner
from optuna.trial._state import TrialState


class LastPlacePruner(BasePruner):
    def __init__(self, warmup_steps, warmup_trials):
        self._warmup_steps = warmup_steps
        self._warmup_trials = warmup_trials

    def prune(self, study: "optuna.study.Study", trial: "optuna.trial.FrozenTrial") -> bool:
        # Get the latest score reported from this trial
        step = trial.last_step

        if step:  # trial.last_step == None when no scores have been reported yet
            this_score = trial.intermediate_values[step]

            # Get scores from other trials in the study reported at the same step
            completed_trials = study.get_trials(deepcopy=False, states=(TrialState.COMPLETE,))
            other_scores = [
                t.intermediate_values[step]
                for t in completed_trials
                if step in t.intermediate_values
            ]
            other_scores = sorted(other_scores)

            # Prune if this trial at this step has a lower value than all completed trials
            # at the same step. Note that steps will begin numbering at 0 in the objective
            # function definition below.
            if step >= self._warmup_steps and len(other_scores) > self._warmup_trials:
                if this_score < other_scores[0]:
                    print(f"prune() True: Trial {trial.number}, Step {step}, Score {this_score}")
                    return True

        return False

最后, 让我们通过一个简单的超参数优化来确认该实现是正确的.

def objective(trial):
    iris = load_iris()
    classes = np.unique(iris.target)
    X_train, X_valid, y_train, y_valid = train_test_split(
        iris.data, iris.target, train_size=100, test_size=50, random_state=0
    )

    loss = trial.suggest_categorical("loss", ["hinge", "log", "perceptron"])
    alpha = trial.suggest_float("alpha", 0.00001, 0.001, log=True)
    clf = SGDClassifier(loss=loss, alpha=alpha, random_state=0)
    score = 0

    for step in range(0, 5):
        clf.partial_fit(X_train, y_train, classes=classes)
        score = clf.score(X_valid, y_valid)

        trial.report(score, step)

        if trial.should_prune():
            raise optuna.TrialPruned()

    return score


pruner = LastPlacePruner(warmup_steps=1, warmup_trials=5)
study = optuna.create_study(direction="maximize", pruner=pruner)
study.optimize(objective, n_trials=50)

Out:

prune() True: Trial 6, Step 1, Score 0.72
prune() True: Trial 12, Step 1, Score 0.62
prune() True: Trial 13, Step 1, Score 0.66
prune() True: Trial 14, Step 1, Score 0.7
prune() True: Trial 16, Step 1, Score 0.5
prune() True: Trial 17, Step 3, Score 0.5
prune() True: Trial 19, Step 4, Score 0.7
prune() True: Trial 30, Step 1, Score 0.8
prune() True: Trial 43, Step 4, Score 0.64
prune() True: Trial 49, Step 4, Score 0.64