Optuna Tutorial with Pytorch

先日PFNからハイパーパラメータチューニングを自動でやってくれるというフレームワークが公開されました。

optuna.org

PFN内でもOpen Images Challenge 2018の際にはこれを用いてパラメータチューニングをしていたとか。

これは使うっきゃない！！

ということで、PytorchでMNISTを通してoptunaのtutorialをやります！

最近流行りのgoogle colaborator上で試してみました。

ちなみに2019年2月からgoogle colabでpytorchがデフォルトインストールされたようです。

環境

google colabrator (2019/3)

optuna: 0.9.0
pytorch: 1.0.1.post2

1. Optunaのインストール

インストール方法は公式を参考に。

pip install optuna

google colab上では先頭に!をつける必要があります。

!pip install optuna

2. MNISTデータの準備

今回は簡単のためMNISTデータを使って遊んでいきます。

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


# バッチサイズ
BATCH_SIZE = 64

# 入力画像を[-1, 1]に正規化
# 参考：https://discuss.pytorch.org/t/understanding-transform-normalize/21730
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])

# MNISTデータのダウンロード&DataLoaderインスタンス作成
train_set = MNIST(root='./data', train=True, download=True, transform=transform)
train_loader = DataLoader(train_set, batch_size=BATCH_SIZE, shuffle=True, num_workers=2)
test_set = MNIST(root='./data', train=False,  download=True, transform=transform)
test_loader = DataLoader(test_set, batch_size=BATCH_SIZE, shuffle=False, num_workers=2)

3. モデルの設計

まだOptunaは出てきません。

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

class Net(nn.Module):
    def __init__(self, activation):
        super(Net, self).__init__()
        self.activation = activation
        self.conv1 = nn.Conv2d(1, 16, kernel_size=3)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3)
        self.conv3 = nn.Conv2d(32, 10, kernel_size=3)

    def forward(self, x):
        x = self.activation(F.max_pool2d(self.conv1(x), 2))
        x = self.activation(F.max_pool2d(self.conv2(x), 2))
        x = self.activation(F.max_pool2d(self.conv3(x), 2))
        x = F.adaptive_avg_pool2d(x, output_size=(1, 1))
        return x.view(-1, 10)

4. トレーニング & テスト用設定

def train(model, device, train_loader, optimizer, criterion):
    model.train()
    for data, target in train_loader:
      data, target = data.to(device), target.to(device)
      
      # Zero the parameter gradients
      optimizer.zero_grad()
      
      # forward + backward + optimize
      output = model(data)
      
      loss = criterion(output, target)
      loss.backward()
      
      optimizer.step()
      

def evaluate(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)
            _, predicted = torch.max(output.data, 1)
            correct += (predicted == target).sum().item()
    accuracy = 1 - correct / len(test_loader.dataset)
    return accuracy

5. Optunaを用いたハイパーパラメータの設定

ようやくOptunaの実装が始まります。

記事執筆現在では、Optunaでは5つのパラメータ設定方法があります。

# 1. Categorical parameter
optimizer = trial.suggest_categorical('optimizer', ['MomentumSGD', 'Adam'])

# 2. Int parameter
num_layers = trial.suggest_int('num_layers', 1, 3)

# 3. Uniform parameter
dropout_rate = trial.suggest_uniform('dropout_rate', 0.0, 1.0)

# 4. Loguniform parameter
learning_rate = trial.suggest_loguniform('learning_rate', 1e-5, 1e-2)

# 5. Discrete-uniform parameter
drop_path_rate = trial.suggest_discrete_uniform('drop_path_rate', 0.0, 1.0, 0.1)

上記の例のようにハイパーパラメータそれぞれの特性に合わせてサンプリング方法を選択し、探索したいパラメータを設定します。

今回は、optimizer, activation, weight_decay, learning rateをOptunaを使って最適化できるようにしていきます。

import torch.optim as optim

def get_optimizer(trial, model):
  # optimizer をAdamとMomentum SGDで探索
  optimizer_names = ['Adam', 'MomentumSGD']
  optimizer_name = trial.suggest_categorical('optimizer', optimizer_names)
    
  # weight decayの探索
  weight_decay = trial.suggest_loguniform('weight_decay', 1e-10, 1e-3)
  
  # optimizer_nameで分岐
  if optimizer_name == optimizer_names[0]: 
      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


def get_activation(trial):
  # 活性化関数の探索. ReLU or ELu.
    return trial.suggest_categorical('activation', [F.relu, F.elu])

6. 目的関数の設計

Optunaは現在、目的関数の最小化のみをサポートしているため目的関数の返り値は小さくなることが目標になるように設計する必要があります。そのため目的関数の返り値をaccuracyではなく、error rateに設定しています。

import optuna
import torch

EPOCH = 5

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

    activation = get_activation(trial)

    model = Net(activation).to(device)
    optimizer = get_optimizer(trial, model)
    criterion = nn.CrossEntropyLoss()
    
    # Training
    for step in range(EPOCH):
        train(model, device, train_loader, optimizer, criterion)
        
    # Evaluation
    accuracy = evaluate(model, device, test_loader)
  
    # 返り値が最小となるようにハイパーパラメータチューニングが実行される
    return 1.0 - accuracy

7. Optunaの実行

ここまでで準備ができたため、次は実際に最適化を実行していきます。

まずoptuna.create_study()で学習用のインスタンスを作成します。そしてoptimize()メソッドで目的関数の返り値を元にパラメータチューニングを実行します。

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

8. 結果の確認

以下の操作で各種結果が確認できます。

tudy.best_params  # Get best parameters for the objective function.
study.best_value  # Get best objective value.
study.best_trial  # Get best trial's information.
study.trials  # Get all trials' information.
study.trials_dataframe()  # Get a pandas dataframe like

番外編

RDB(Relational Data Base)を利用した最適化結果の保存

簡単にSQLiteを利用した結果の保存が可能になっています。

study_name = 'example-study'  # Unique identifier of the study.
study = optuna.create_study(study_name=study_name, storage='sqlite:///example.db')

最適化の再開

create_study()の引数であるload_if_exsitsをTrueにすることで保存した結果を読み込んで最適化を再開可能です。

study = optuna.create_study(study_name='example-study', storage='sqlite:///example.db', load_if_exists=True)
study.optimize(objective, n_trials=3)

attributeの追加

set_usr_attr()を用いて任意の値を結果に保存していくことができます。

def objective(trial):
  accuracy = ...
  trial.set_user_attr('accuracy', accuracy)
  return 1.0 - accuracy

まとめ

元々のコードをほとんどいじることなくパラメータチューニングを行えるため非常に便利ですね。

さすがPFN。俺達にできないことを平然とやってのけるッ

そこにシビれる！あこがれるゥ！

0.5から始める機械学習

Machine Learning, Deep Learning, Computer Vision に関する備忘録

【Optuna】Optuna Tutorial with Pytorch