Sansan Tech Blog

Sansanのものづくりを支えるメンバーの技術やデザイン、プロダクトマネジメントの情報を発信

【Zoom or Die】第2回 Hydra+Axでハイパーパラメータサーチ

機械学習 Python R&D 連載

f:id:S_aiueo321:20200529093604p:plain

こんにちは，DSOC 研究開発部の内田です． 最近すっかり秋めいてきて，短パン小僧 の私としてはちょっと寒いくらいです． 涼しくなってきたので1人で寺巡りをすることが多いのですが．そのせいで夏本番より日焼けして短パン小僧感が増す始末です．

先日同研究員の高橋が寄稿した記事でもMLOpsに入門していましたが，最近社内になんとなくMLOpsの風を感じます． そんなわけで今回はHydra+Axを使ってハイパーパラメータサーチをしてみたいと思います．

buildersbox.corp-sansan.com

ツールの紹介

Hydra

機械学習モデル，特に深層なモデルともなれば膨大な数のハイパーパラメータを有します． これらを調整・管理する際，皆さんはどうされているでしょうか．

Pythonを利用している場合， argparse を介してコマンドライン引数から調整するのが手っ取り早い気がします． しかしながら，コマンドラインからパラメータを漏れなく指定するのは意外と骨が折れます． そこで，静的に設定ファイルを作成しておき，それらを読み込んで学習・テストを回す方法があります．

Hydraは，Facebook Open Sourceから公開されている設定管理ツールです． YAML形式でパラメータを保持しておき，@hydra.main() デコレータを介してそれらを取得することが可能です． 設定ファイルの切替やパラメータのオーバーライドに長けていて，私自身ここ半年くらい色んなところで利用しています．

hydra.cc

Ax

膨大な数のハイパーパラメータを手動で調整するのは辛いものです． グリッドサーチなどで自動化することも可能ですが，探索範囲が広くなるにつれて試行回数が現実的ではなくなっていきます．

この問題を解決するため，最近ではベイズ最適化を利用してパラメータ探索を自動化するライブラリがいくつかあります． 代表的な物を挙げれば，Preferred Networksが公開しているOptunaが有名でしょうか． これらのライブラリとHydraを組み合わせることで，効率的なハイパーパラメータサーチが実現できると考えます．

Ax(Adaptive Experimentation Platform)はHydraと同様，Facebook Open Sourceが公開しているパラメータサーチを自動化するライブラリです． Axはガウス過程を利用して最適化を行うことができ，内部的にBoTorchが動作しています． 公開元が同じということもあってか，AxはHydraからプラグインとして利用できるので親和性が高いです．

ax.dev

実験

ここでは，MNIST*1に対する識別モデルのハイパーパラメータをHydra+Axで最適化したいと思います．

実験環境

今回の実験ではPython3.7を利用します．主な依存パッケージは下記です．

hydra = "^2.5"
hydra-ax-sweeper = "^1.0.0-rc.6"
ax-platform = "^0.1.14"
torch = "^1.6.0"
torchvision = "^0.7.0"

ファイルの配置はこんな感じです．

.
├── config
│   └── train.yaml
└── train.py

やったこと

モデル構築

今回識別モデルには，畳み込み層が2つ，全結合層が3つの畳み込みニューラルネットを使用します． コンストラクタの引数に渡している cfg はHydra経由で取得したハイパーパラメータの辞書*2です． コードからわかるように，畳み込みのチャネル数 n_channels と全結合層の次元数 n_hidden が設定ファイルから指定可能です．

import torch.nn as nn
from omegaconf import DictConf


class Net(nn.Sequential):
    def __init__(self, cfg: DictConfig) -> None:
        super().__init__(
            nn.Conv2d(1, cfg.n_channels, kernel_size=3, padding=1),
            nn.ReLU(True),
            nn.MaxPool2d(2, 2),

            nn.Conv2d(cfg.n_channels, cfg.n_channels * 2, kernel_size=3, padding=1),
            nn.ReLU(True),
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),

            nn.Linear(cfg.n_channels * 2, cfg.n_hidden),
            nn.ReLU(True),
            nn.Linear(cfg.n_hidden, cfg.n_hidden // 2),
            nn.ReLU(True),
            nn.Linear(cfg.n_hidden // 2, 10),
            nn.Softmax(dim=1)
        )

学習コード

学習コードはこんな感じです．モデルに加えて，学習率と最大エポック数に関しても設定ファイルから指定可能としています． train() の中身は長くなるので割愛しますが，1エポック毎に検証データに対する誤差を記録し，最終的に最も良い検証誤差を返却します． その際，10エポック改善がない場合早期終了します． Axプラグインを利用するには，@hydra.main() デコレータが付いたメソッドの返り値に学習の最終的なスコアを指定する必要があります．

import multiprocessing as mp

import hydra
import numpy as np
import torch.optim as optim
import torchvisioin.transforms as transforms
from torch.utils.data import DataLoader
from torchvisioin.datasets import MNIST


@hydra.main(config_path='config', config_name='train')
def app(cfg: DictConfig) -> float:
    train_dataset = MNIST(
        root=hydra.utils.get_original_cwd() + 'data',
        download=True,
        train=True,
        transform=transforms.Compose([lambda x: np.array(x), transforms.ToTensor()]),
        target_transform=lambda x: torch.as_tensor(x)
    )
    val_dataset = MNIST(
        root=hydra.utils.get_original_cwd() + 'data',
        download=True,
        train=False,
        transform=transforms.Compose([lambda x: np.array(x), transforms.ToTensor()]),
        target_transform=lambda x: torch.as_tensor(x)
    )

    train_dataloader = DataLoader(
        train_dataset,
        batch_size=cfg.batch_size,
        num_workers=mp.cpu_count(),
        pin_memory=True,
        shuffle=True,
        drop_last=True
    )
    val_dataloader = DataLoader(
        val_dataset,
        batch_size=cfg.batch_size,
        num_workers=mp.cpu_count(),
        pin_memory=True,
        drop_last=True
    )

    model = Net(cfg).to(DEVICE)
    optimizer = optim.Adam(model.parameters(), lr=cfg.lr)

    result: float = train(model, optimizer, train_dataloader, val_dataloader, cfg.max_epochs)

    return result


if __name__ == "__main__":
    app()

設定ファイル

最後に設定ファイルを記述します． 設定ファイルはハイパーパラメータの設定とAxの設定の大きく2つに分かれています． 前者にはコード内でcfgから取得するパラメータを記述します． 後者については，まずAxプラグインを利用することを宣言し，細かい設定を行なっていきます． 設定する項目としては，最適化の最大試行回数や方向，最適化対象のパラメータがあります． 今回は最大試行回数を5，返却されるスコアが誤差であるため最適化の方向は最小化とし，n_channels，n_hidden，lrを最適化します． 値の選択方法には3種類あり，指定の候補から選択するchoice，指定の範囲から値をサンプルするrange，特定の値を表すfixedがあります． params 以下の変数名にアンダースコアを使いたい場合は，エスケープが必要ということに注意してください*3．

# ハイパーパラメータ
n_channels: 32
n_hidden: 512
lr: 0.001
max_epochs: 100

# Axプラグインを利用することを宣言
defaults:
  - hydra/sweeper: ax

# Axの設定
hydra:
  sweeper:
    ax_config:
      max_trials: 5  # 最大試行回数
      experiment:
        minimize: true  # 最適化の方向

      # 最適化対象のハイパーパラメータ
      params:
        n\\_channels:
          type: choice
          values: [16, 32, 64]
        n\\_hidden:
          type: choice
          values: [256, 512, 1024]
        lr:
          type: range
          bounds: [0.0001, 0.001]

結果

上記のコードを用いて，実際にハイパーパラメータを最適化してみます．実行コマンドは python main.py -m です．

[2020-09-17 16:04:14,870][HYDRA] AxSweeper is optimizing the following parameters: 
n_channels: choice=[16, 32, 64], type = int
n_hidden: choice=[256, 512, 1024], type = int
lr: range=[0.0001, 0.001], type = float
[INFO 09-17 16:04:14] ax.modelbridge.dispatch_utils: Using Sobol generation strategy.
[2020-09-17 16:04:14,951][HYDRA] AxSweeper is launching 5 jobs
[2020-09-17 16:04:15,975][HYDRA] Launching 5 jobs locally
[2020-09-17 16:04:15,976][HYDRA]        #0 : lr=0.0007630527009256184 n_channels=32 n_hidden=256
epoch:99, best_score: 1.4866631581233098, patience: 1: 100%|██| 100/100 [05:12<00:00,  3.13s/it]
[2020-09-17 16:09:32,368][HYDRA]        #1 : lr=0.0008976607835851609 n_channels=16 n_hidden=1024
epoch:99, best_score: 1.4851646851270626, patience: 3: 100%|██| 100/100 [05:11<00:00,  3.11s/it]
[2020-09-17 16:14:44,080][HYDRA]        #2 : lr=0.00022952019199728963 n_channels=64 n_hidden=1024
epoch:99, best_score: 1.486624574049925, patience: 6: 100%|██| 100/100 [06:15<00:00,  3.76s/it]
[2020-09-17 16:21:00,157][HYDRA]        #3 : lr=0.0007424576037563384 n_channels=16 n_hidden=512
epoch:99, best_score: 1.4902324921045549, patience: 3: 100%|██| 100/100 [05:12<00:00,  3.12s/it]
[2020-09-17 16:26:12,558][HYDRA]        #4 : lr=0.0009297236477956176 n_channels=32 n_hidden=256
epoch:92, best_score: 1.4843894640604656, patience: 10:  92%|██ | 92/100 [05:05<00:26,  3.32s/it]
[2020-09-17 16:31:18,438][HYDRA] New best value: 1.4843894640604656, best parameters: {'lr': 0.0009297236477956176, 'n_channels': 32, 'n_hidden': 256}
[2020-09-17 16:31:18,440][HYDRA] Best parameters: {'lr': 0.0009297236477956176, 'n_channels': 32, 'n_hidden': 256}

試行回数がたった5回ですので結果の信頼性には目を瞑るとして，無事 Best parameters: {'lr': 0.0009297236477956176, 'n_channels': 32, 'n_hidden': 256} という値を得ることができました． ついでにTensorBoardで学習曲線を出してみました．なんかそれっぽいグラフが出てきていますね．

f:id:S_aiueo321:20200917180929p:plain
学習曲線

まとめ

今回はHydra+Axを使ってハイパーパラメータサーチをやってみました． コードをほぼ変更することなく，非常に簡単に設定ファイルからハイパーパラメータサーチができるのがよかったです！

ただ，私の調査不足かも知れませんが，今回の方法だと学習の枝刈りができないため，試行回数分の学習がフルで回るor古典的な早期終了しかできないという状況です． 一応設定項目にはearly_stopがあるのですが，利用方法がわからず… どなたか知見があれば教えていただきたいです． その点Optunaは枝刈りをサポートしているので，学習時間などを考慮するとHydra+Optunaが使いやすいかも知れませんね*4． 両者とも，今後の業務で利用してみて，感想があればどこかで共有したいと思います！


▼これまでの記事

【Zoom or Die】第1回 NTIRE2020 Challenge 結果速報

*1:LeCun, Yann. "The MNIST database of handwritten digits." http://yann.lecun.com/exdb/mnist/ (1998).

*2:実際には組み込みの dict を拡張した DIctConfig が取得できます．

*3:エスケープしない場合，ドットに置き換わってしまいます．

*4:Optunaを利用する場合は，objectiveメソッド内でcfgを書き換える処理を自前で書けば，今回とほぼ同様の設定ファイル変更でハイパーパラメータサーチができると思っています．

© Sansan, Inc.