機械学習
DeepLearning
PyTorch
Optuna

はじめに

先日(12月3日)、PFNからOptunaという、機械学習向けハイパーパラメータ自動最適化フレームワークが公開されました

本記事では、OptunaをPyTorchと一緒に使って、ハイパーパラメータチューニングを行います。

ソースコード:https://github.com/eduidl/optuna-with-pytorch/blob/master/mnist.ipynb (色々やり直したりしたので結果が変わっています。)

環境

  • Python 3.6.7
  • PyTorch 0.4.1
  • Optuna 0.4.0

インストール

pip install optuna

とするだけです。

パラメータチューニング

公式リポジトリにChainerを使ったサンプルがあるので、それを参考にしました。

ここでは、

  • 活性化関数 (ReLU or ELU)
  • 最適化手法 (Adam or MomentumSGD)
  • 学習率
  • weight_decay
  • ドロップアウトの選出確率

の5つのパラメータを振ることにしています(AdamとMomentumSGDで学習率のパラメータ名が異なるので、6つともいえるが、実質5つ)。

データセットの準備

本筋ではないですが、一応書いておきます。

from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
from torchvision import transforms
import numpy as np

BATCHSIZE = 128

transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.5,), (0.5,))])

train_set = MNIST(root='./data', train=True,
                  download=True, transform=transform)
train_loader = DataLoader(train_set, batch_size=BATCHSIZE,
                          shuffle=True, num_workers=2)

test_set = MNIST(root='./data', train=False, 
                 download=True, transform=transform)
test_loader = DataLoader(test_set, batch_size=BATCHSIZE,
                         shuffle=False, num_workers=2)

classes = tuple(np.linspace(0, 9, 10, dtype=np.uint8))

モデル

PyTorchの公式サンプルコードを参考にしつつ、活性化関数と、ドロップアウトの選出確率を外から渡せるようにしています。
層の深さのような"ダイナミック"なパラメータを変更する方が面白いのでしょうが、手を抜いてしていません。

import torch
import torch.nn as nn
import torch.nn.functional as F

class Net(nn.Module):
    def __init__(self, activation, dropout_prob):
        super(Net, self).__init__()
        self.activation = activation
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.conv2_drop = nn.Dropout2d(p=dropout_prob)
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)

    def forward(self, x):
        x = self.activation(F.max_pool2d(self.conv1(x), 2))
        x = self.activation(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        x = x.view(-1, 320)
        x = self.activation(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)

TrainとTest

これもあまり本筋ではないですが、一応載せておきます。

def train(model, device, train_loader, optimizer):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()

def test(model, device, test_loader):
    model.eval()
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            pred = output.max(1, keepdim=True)[1]
            correct += pred.eq(target.view_as(pred)).sum().item()

    return 1 - correct / len(test_loader.dataset)

Optimizer

Optunaの公式サンプルコードを多少変更すると以下のようになります。

import torch.optim as optim

def get_optimizer(trial, model):
    optimizer_names = ['Adam', 'MomentumSGD']
    optimizer_name = trial.suggest_categorical('optimizer', optimizer_names)
    weight_decay = trial.suggest_loguniform('weight_decay', 1e-10, 1e-3)
    if optimizer_name == optimizer_names[0]: 
    # サンプルコードでは
    # if optimizer_name == 'Adam'
    # となっているが、地味にバグの元になりそうな気がした
        adam_lr = trial.suggest_loguniform('adam_lr', 1e-5, 1e-1)
        optimizer = optim.Adam(model.parameters(), lr=adam_lr, weight_decay=weight_decay)
    else:
        momentum_sgd_lr = trial.suggest_loguniform('momentum_sgd_lr', 1e-5, 1e-1)
        optimizer = optim.SGD(model.parameters(), lr=momentum_sgd_lr,
                              momentum=0.9, weight_decay=weight_decay)
    return optimizer

やっとOptunaらしいところが出てきました。

lists = ['Adam', 'MomentumSGD']
optimizer_name = trial.suggest_categorical('optimizer', lists)

trial.suggest_categorical は、Enumであるようなパラメータを定義します。
今回でいうと、optimzer_nameAdamMomentumSGDになるということです。

weight_decay = trial.suggest_loguniform('weight_decay', 1e-10, 1e-3)

trial.suggest_loguniform は、例えば学習率を$0.1, 0.03, 0.01,...$と振っていくように、値を$x$倍して変更していくようなパラメータを定義します。adam_lr, momentum_sgd_lrに関しては同様なので割愛。

get_optimizer の改善?

上のコードをみて、suggest_loguniformに比べ、suggest_categoricalの書き方がいけていないと思いませんか?私はそう思います。

ところで、私はPythonという言語があまり好きではない(変数のスコープとか)のですが、関数がファーストクラスオブジェクトであるところは結構好きです。

急に言語Disを持ち出して何が言いたいかというと、suggest_categoricalに関数のlistを渡せばよいのではという話です。

def adam(model, trial, weight_decay):
    adam_lr = trial.suggest_loguniform('adam_lr', 1e-5, 1e-1)
    return optim.Adam(model.parameters(), lr=adam_lr, weight_decay=weight_decay)

def momentum(model, trial, weight_decay):
    momentum_sgd_lr = trial.suggest_loguniform('momentum_sgd_lr', 1e-5, 1e-1)
    return optim.SGD(model.parameters(), lr=momentum_sgd_lr,
                     momentum=0.9, weight_decay=weight_decay)

def get_optimizer(trial, model):
    optimizer = trial.suggest_categorical('optimizer', [adam, momentum])
    weight_decay = trial.suggest_loguniform('weight_decay', 1e-10, 1e-3)
    return optimizer(model, trial, weight_decay)

こちらの方が綺麗ではないでしょうか(※個人的感想です)。ただ、実際に動作が保証されているのかは自信ありません。
(手元ではぱっと見動いているし、study.trials.params_in_internal_reprの中身からすると、Enumよろしく実体は整数っぽいし、特に問題はないと思うのですが。)

活性化関数

Optimizerと同様です。

def get_activation(trial):
    activation_names = ['ReLU', 'ELU']
    activation_name = trial.suggest_categorical('activation', activation_names)
    if activation_name == activation_names[0]:
        activation = F.relu
    else:
        activation = F.elu
    return activation

また、suggest_categoricalに関数オブジェクトを渡してよいなら、

def get_activation(trial):
    return trial.suggest_categorical('activation', [F.relu, F.elu])

と書けるはずです。

メイン

必要な関数が出揃ったのでようやく、パラメータチューニングに入れます。

import optuna

EPOCH = 10

def objective(trial):
    device = "cuda" if torch.cuda.is_available() else "cpu"

    activation = get_activation(trial)
    dropout_prob = trial.suggest_uniform("dropout_prob", 0, 0.8)

    model = Net(activation, dropout_prob).to(device)
    optimizer = get_optimizer(trial, model)

    for step in range(EPOCH):
        train(model, device, train_loader, optimizer)
        error_rate = test(model, device, test_loader)

    return error_rate

study = optuna.create_study()
study.optimize(objective, n_trials=100)

Optunaでパラメータチューニングを行う場合、まずoptuna.create_study()optuna.study.Studyクラスのインスタンスを生成し、インスタンスメソッドoptimze()を呼ぶことにより行えますが、その際に引数として目的関数を与える必要があります。
現在のOptunaでは、目的関数の返り値が小さくなるようにパラメータ探索を行います。1そのため、目的関数であるobjective関数は、正解率ではなく、誤り率を返すようにしています。

また、n_trainsで何パターン試すかを指定しています。

結果

study.optimize(objective, n_trials=100) とすると以下のようにログが表示されていきます。現在の目的関数の値、目的関数の最小値、最小値を達成したパラメータが表示されます。

一応、https://optuna.readthedocs.io/en/latest/reference/logging.html を参考にログレベルを調整すれば、ログを吐かないようにもできるようですが、一切ログがでないのもそれはそれで辛いので、私はログを吐かせました。
(tqdmを使えたらいいんですが、n_trials分のループは隠蔽されていて使えない気がします。間違いでした。追記2に書いた方法でできそうです。)

[I 2018-12-05 06:12:41,057] Finished a trial resulted in value: 0.09560000000000002. Current best value is 0.09560000000000002 with parameters: {'activation': 'ReLU', 'dropout_prob': 0.35100773000855123, 'optimizer': 'Adam', 'weight_decay': 6.227514717616339e-06, 'adam_lr': 1.541332256542363e-05}.
[I 2018-12-05 06:13:06,201] Finished a trial resulted in value: 0.35619999999999996. Current best value is 0.09560000000000002 with parameters: {'activation': 'ReLU', 'dropout_prob': 0.35100773000855123, 'optimizer': 'Adam', 'weight_decay': 6.227514717616339e-06, 'adam_lr': 1.541332256542363e-05}.
[I 2018-12-05 06:13:31,358] Finished a trial resulted in value: 0.01849999999999996. Current best value is 0.01849999999999996 with parameters: {'activation': 'ReLU', 'dropout_prob': 0.675273078348647, 'optimizer': 'Adam', 'weight_decay': 1.3423197864054106e-09, 'adam_lr': 0.0031045501711910307}.
[I 2018-12-05 06:13:56,193] Finished a trial resulted in value: 0.011399999999999966. Current best value is 0.011399999999999966 with parameters: {'activation': 'ELU', 'dropout_prob': 0.43004091975326153, 'optimizer': 'Adam', 'weight_decay': 8.600682467290723e-08, 'adam_lr': 0.0006490359130172907}.
[I 2018-12-05 06:14:21,261] Finished a trial resulted in value: 0.04390000000000005. Current best value is 0.011399999999999966 with parameters: {'activation': 'ELU', 'dropout_prob': 0.43004091975326153, 'optimizer': 'Adam', 'weight_decay': 8.600682467290723e-08, 'adam_lr': 0.0006490359130172907}.
....

終了後、以下のコマンドで結果が確認できます。

> study.best_params
{'activation': 'ReLU',
 'dropout_prob': 0.1750391800071884,
 'optimizer': 'MomentumSGD',
 'weight_decay': 1.185808055634304e-07,
 'momentum_sgd_lr': 0.02022072772568203}

> study.best_value
0.006900000000000017

> df = study.trials_dataframe() # pandasのDataFrame形式
> df.head()

image.png

選ばれたパラメータは以下の通りで、

activation ReLU
dropout_prob 0.35100773000855123
optimizer Adam
weight_decay 6.227514717616339e-06
adam_lr 1.541332256542363e-05

accuracyは99.3%くらいでした。

(枝刈り[参考]もやりましたが、書くのに疲れたので終わりとします。)

追記

よく考えると、コンストラクタにtrailを渡して下のように書く方が良かったですね。

class Net(nn.Module):
    def __init__(self, trial):
        super(Net, self).__init__()
        self.activation = get_activation(trial)
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.conv2_drop = nn.Dropout2d(p=trial.suggest_uniform("dropout_prob", 0, 0.8))
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)
    # 中略

def objective(trial):
    device = "cuda" if torch.cuda.is_available() else "cpu"

    model = Net(activation, dropout_prob).to(device)
    optimizer = get_optimizer(trial, model)

    for step in range(EPOCH):
        train(model, device, train_loader, optimizer)
        error_rate = test(model, device, test_loader)

    return error_rate

追記2

tqdmと組み合わせる方法も思いつきました。
下のようにラッパーを作って、外からtqdmオブジェクトを渡して、 update() を呼ぶようにしてやればよいですね。
(枝刈りをする場合は、raise optuna.structs.TrialPruned() をする前にも、pbar.update() しましょう。(一敗))

# jupyter notebook
from tqdm import tqdm_notebook as tqdm

# optunaのログを無効にする
optuna.logging.disable_default_handler()

def objective_wrapper(pbar):
    def objective(trial):
        device = "cuda" if torch.cuda.is_available() else "cpu"

        model = Net(trial).to(device)
        optimizer = get_optimizer(trial, model)

        for step in range(EPOCH):
            train(model, device, train_loader, optimizer)
            error_rate = test(model, device, test_loader)

        pbar.update()

        return error_rate

    return objective

TRIAL_SIZE = 100
with tqdm(total=TRIAL_SIZE) as pbar:
    study = optuna.create_study()
    study.optimize(objective_wrapper(pbar), n_trials=TRIAL_SIZE)

image.png

参考:https://pypi.org/project/tqdm/#manual

(どうでもいい)疑問

そういえば、Optuna = Optimize + Unagi という説を見かけたが事実なんだろうか?

参考リンク


  1. 将来的には、create_study(direction='maximize')とすれば大きくすることを目指すようにすることもできるようになると思われるが、少なくとも2018/12/5現在はドキュメント によると Note that maximize is currently unsupported とのことです。(線型計画ソルバだと、最小値の代わりに-1倍したもの最大値を計算するとかしていた気がするが、そう簡単な話ではないのか?)