Optuna的进阶使用

Optuna作为主要面向深度学习超参数调优开发的框架，在实现之初就考虑到了大型模型参数调优的各种实际情况，并逐一针对它们设计了解决方案。

分布式优化

一提到分布式算法，我们想到的可能是麻烦的 debug 过程和分布式实现过程中各种线程锁之类的问题。你可能会好奇，带有 GIL的 python 超参数优化库是怎么实现分布式优化的？实际上，通过选择不同的共享参数的方式，Optuna 规避了这一问题（optuna 在内部使用了 joblib, 然而仍然受到 GIL 的限制）。在 Optuna 中，用户不是在单个脚本内派生出不同的线程，而是通过针对同一个 study 启动不同的优化过程来实现分布式优化的。

在一个分布式的 study 优化过程中，用户首先创建一个空 study:

$ optuna create-study --study-name "distributed-example" --storage "sqlite:///example.db"

\[I 2018\-10-31 18:21:57,885\] A new study created with name: distributed-example

这个命令会在当前目录下生成一个数据库文件

$ ls
example.db foo.py
(doc-optuna)

该 study 将会用于后期的参数和历史记录存储。一个中心化的参数存储是有必要的，因为在超参数优化过程中，一个后续的试验（trial）的参数采样范围会受到前面的参数采样历史记录影响。而由于Optuna 透过数据库接口来读写这些数据，因此多个 client 之间不会造成干扰，用户只需在不同的进程中运行这个 study 即可：比如在不同的终端窗口中运行优化脚本。

假如优化脚本如下：

import optuna

def objective(trial):
    x = trial.suggest_uniform('x', -10, 10)
    return (x - 2) ** 2

if __name__ == '__main__':
    study = optuna.load_study(study_name='distributed-example', storage='sqlite:///example.db')
    study.optimize(objective, n_trials=100)

（注意，这里的 n_trials 不是指 对于该 study 总共会运行100次试验，而是每一个单独的优化进程都会跑100次试验，n个进程下的总试验次数是 n x 100)

考虑两个进程的优化，我们可以在同一个目录下的两个终端窗口中运行:

$ python foo.py
[I 2018-10-31 18:46:44,308] Finished a trial resulted in value: 1.1097007755908204. Current best value is 0.00020881104123229936 with parameters: {'x': 2.014450295541348}.
[I 2018-10-31 18:46:44,361] Finished a trial resulted in value: 0.5186699439824186. Current best value is 0.00020881104123229936 with parameters: {'x': 2.014450295541348}.
...

等优化完成以后开启第三个脚本检查总的trial数，可以看到总试验次数是200：

>>> len(study.trials)
200

此时，如果打开数据库，我们会发现，试验和试验所用到的参数是分开存储的：红色是试验记录，蓝色是参数采样记录。通过分隔着两种记录，不同的进程从同一个空间中采样变得更加容易。

这种分布式优化的办法看起来似乎有些“手动”，然而却是可扩展性最佳的做法之一。因为，你甚至可以将存储优化历史记录和参数采样的 study 设置成远程数据库存储。这样的话，即使要在不同的机器上同时对同一个模型进行调参也非常方便：拷贝一份模型代码，然后设置好数据库，启动优化进程即可。

在跑优化时还有一点需要注意：在大型的分布式超参优化过程中，最好不要用 SQLite 这种数据库，使用成熟的 MySQL 或者 PostgreSQL 是可靠的选择。

剪枝算法

在复杂模型的优化过程中，即使在训练的早期我们也常常能看到有一些参数组合是明显无效的，此时如果对该参数组合继续进行优化毫无意义，对其进行提前终止能节省大量的运算时间，而各种剪枝算法也应运而生。用户当然可以自行实现这些算法，但是这通常要求对现存的模型代码进行大量的修改，如果剪枝的标准涉及优化过程历史记录比较的话（比如对中位数以下进行剪枝），实现起来则更加麻烦。

Optuna 为了解决这一问题，通过提供 Pruner 接口，最大程度地将剪枝的逻辑和模型逻辑进行了解耦。考虑一个迭代式训练的过程，在每一步之后，你只需调用 report 和 should_prune 函数即可进行剪枝。其中 report 负责监测迭代过程中的目标函数中间值，而 should_prune 则进行判断，选择是否触发一个异常，让 Optuna终止该试验。

下面是一个来自官方文档的剪枝例子：

import sklearn.datasets
import sklearn.linear_model
import sklearn.model_selection

import optuna

def objective(trial):
    iris = sklearn.datasets.load_iris()
    classes = list(set(iris.target))
    train_x, valid_x, train_y, valid_y = \
    sklearn.model_selection.train_test_split(iris.data, iris.target, test_size=0.25, random_state=0)

    alpha = trial.suggest_loguniform('alpha', 1e-5, 1e-1)
    clf = sklearn.linear_model.SGDClassifier(alpha=alpha)

    for step in range(100):
        clf.partial_fit(train_x, train_y, classes=classes)

        # Report intermediate objective value.
        intermediate_value = 1.0 - clf.score(valid_x, valid_y)
        trial.report(intermediate_value, step)

    # Handle pruning based on the intermediate value.
    if trial.should_prune():
        raise optuna.TrialPruned()

    return 1.0 - clf.score(valid_x, valid_y)

# Set up the median stopping rule as the pruning condition.
study = optuna.create_study(pruner=optuna.pruners.MedianPruner())
study.optimize(objective, n_trials=20)

不难看出，整个剪枝算法的实现从目标函数中被抽离出来了，通过一个 Pruner 实例，每个trial 只需发送中间目标函数值进行判断即可实现剪枝，原有的训练代码几乎无需改变。

该 study 的试验记录输出如下：

$ python pr.py
[I 2020-07-16 13:25:10,428] Trial 0 finished with value: 0.10526315789473684 and parameters: {'alpha': 9.761488546865748e-05}. Best is trial 0 with value: 0.10526315789473684.
[I 2020-07-16 13:25:10,595] Trial 1 finished with value: 0.02631578947368418 and parameters: {'alpha': 3.918322442650925e-05}. Best is trial 1 with value: 0.02631578947368418.
[I 2020-07-16 13:25:10,763] Trial 2 finished with value: 0.2894736842105263 and parameters: {'alpha': 3.561987583888278e-05}. Best is trial 1 with value: 0.02631578947368418.
[I 2020-07-16 13:25:10,942] Trial 3 finished with value: 0.052631578947368474 and parameters: {'alpha': 0.00569807382728555}. Best is trial 1 with value: 0.02631578947368418.
[I 2020-07-16 13:25:11,112] Trial 4 finished with value: 0.07894736842105265 and parameters: {'alpha': 0.00019141260605620358}. Best is trial 1 with value: 0.02631578947368418.
[I 2020-07-16 13:25:11,119] Trial 5 pruned.
[I 2020-07-16 13:25:11,126] Trial 6 pruned.
[I 2020-07-16 13:25:11,316] Trial 7 finished with value: 0.1578947368421053 and parameters: {'alpha': 9.844155535043529e-05}. Best is trial 1 with value: 0.02631578947368418.
[I 2020-07-16 13:25:11,329] Trial 8 pruned.
[I 2020-07-16 13:25:11,334] Trial 9 pruned.
[I 2020-07-16 13:25:11,341] Trial 10 pruned.
[I 2020-07-16 13:25:11,356] Trial 11 pruned.
[I 2020-07-16 13:25:11,546] Trial 12 finished with value: 0.23684210526315785 and parameters: {'alpha': 0.0033635198424680182}. Best is trial 1 with value: 0.02631578947368418.
[I 2020-07-16 13:25:11,731] Trial 13 finished with value: 0.07894736842105265 and parameters: {'alpha': 0.02974537607602637}. 
...

这个例子中采用的是中位数剪枝，实际上Optuna 已内置了多种能满足一般需求的 Pruner，比如百分位数 pruner，hyperband pruner 和针对特性变量监测的 Threshold pruner 等. 关于它们的具体细节请参考 https://zh-cn.optuna.org/reference/pruners.html

集成模块

在有些高度抽象的机器学习库中，比如 keras 或者 Pytorch Ignite，训练过程本身被抽象成了一个函数。此时像上例那样实现剪枝是不可能的。所幸大多数的库在这些方法中都提供了一个钩子接口，可以传入一个检测函数。Optuna 也为这类框架提供了集成的剪枝模块，比如 optuna.integration.PyTorchIgnitePruningHandler 和  optuna.integration.KerasPruningCallback . 这里有一个 Pruning 集成模块的完整列表 介绍 。

下面我们通过一个 PyTorch Ignite 的例子来看一下集成模块是如何在代码中发挥作用的。

def objective(trial):
    # Create a convolutional neural network.
    model = Net(trial)

    device = "cpu"
    if torch.cuda.is_available():
    device = "cuda"
    model.cuda(device)

    optimizer = Adam(model.parameters())
    trainer = create_supervised_trainer(model, optimizer, F.nll_loss, device=device)
    evaluator = create_supervised_evaluator(model, metrics={"accuracy": Accuracy()}, device=device)

    # Register a pruning handler to the evaluator.
    pruning_handler = optuna.integration.PyTorchIgnitePruningHandler(trial, "accuracy", trainer)
    evaluator.add_event_handler(Events.COMPLETED, pruning_handler)

    # Load MNIST dataset.
    train_loader, val_loader = get_data_loaders(TRAIN_BATCH_SIZE, VAL_BATCH_SIZE)

    @trainer.on(Events.EPOCH_COMPLETED)
    def log_results(engine):
    evaluator.run(val_loader)
    validation_acc = evaluator.state.metrics["accuracy"]
    print("Epoch: {} Validation accuracy: {:.2f}".format(engine.state.epoch, validation_acc))

    trainer.run(train_loader, max_epochs=EPOCHS)

    evaluator.run(val_loader)
    return evaluator.state.metrics["accuracy"]

（我们省略了模型定义等部分）

可以看到，上例中 trainer.run 直接实现了训练过程，我们并不能往其中插入一个 should_prune() ，但是透过一个 PyTorchIgnitePruningHandler 实例，我们实现了向 trainer 中注入一个 pruner，之后 evaluator 接受这个pruner作为参数，在每一个 engine 的 run 终止时(Events.COMPLETED) 调用这个 pruning_handler.

不过要注意，并不是每一个机器学习库的集成模块调用方式都是相同的。由于各个库暴露的 API 或者 API风格不一样，相应的集成模块的实现和使用方式也不尽相同。比如在上面的PyTorch Ignite 例子中，实际上你只需将 trial 一并传递到 handler 函数中即可，本来是可以通过往 evaluator.add_event_handler 的handler 参数背后添加额外参数来实现的，但是为了代码的简洁，Optuna 采用的是往 callable 类中提前传递 trial作为类属性绕开这个方法，其核心和之前最简单的 pruning 例子是完全相同的：

def __init__(self, trial, metric, trainer):
    # type: (Trial, str, Engine) -> None

    _imports.check()

    self._trial = trial
    self._metric = metric
    self._trainer = trainer

def __call__(self, engine):
    # type: (Engine) -> None

    score = engine.state.metrics[self._metric]
    self._trial.report(score, self._trainer.state.epoch)
    if self._trial.should_prune():
    message = "Trial was pruned at {} epoch.".format(self._trainer.state.epoch)
    raise optuna.TrialPruned(message)

（ PyTorchIgnitePruningHandler 类内部实现）

可以看到，PyTorchIgnitePruningHandler 只是给 should_prune 做一个包装而已。所以，如果你发现你常用的库目前没有插件但是提供了对应的 callback 接口的话，Optuna 也欢迎你作为开发者来提PR！为一个有callback 的库实现 pruner 并不是一件多么麻烦的事情，却能最大化降低现存代码的修改量。你可以参考 https://github.com/optuna/optuna/tree/master/examples#examples-with-ml-libraries 来找到你常用的机器学习库的对应集成插件列表和应用案例。

其他可定制化的部分

顺便一提，Optuna 的可定制性非常强。比如，采样算法通常会影响参数空间中的搜索效率，如果你发现目前的 sampler 满足不了你的需求的话，你可以自行定制对应的 uder defined sampler. Optuna 提供了一个 BaseSampler, 只要你的新sampler 继承了它并且提供了几个固定的采样接口(sampler_relative 和 sample_independent 等)，就可以在优化中使用任何你想要的参数采样算法而无需改动目标函数代码了。关于Sampler 工作原理的详细介绍，请看 https://zh-cn.optuna.org/reference/samplers.html

类似地，Optuna 还提供了一些更加小巧的定制化方法，比如在 study 和 trial 上都可以使用的 set_user_attr . 这些属性可以让你更便捷地外部修改优化运行时的逻辑而无需改动代码。

Summary

如你所见，Optuna 专为超参数优化而生：方便的目标函数定义，多样化的采样算法和极其便捷的分布式优化让它无愧为大规模调参的首选。目前 Optuna 正处于快速的开发中，但由于1.0 早已发布，其大多数接口都保证了稳定，而许多新特性也将随着 Optuna2.0 的发布一起到来。如果你感觉目前使用的超参数优化方式并不是非常理想的话，不妨尝试一下Optuna。