yacsで実験パラメータを設定する

今回はyacsを用いたパラメータ管理について整理します。

yacs

yacsはPythonコード + YAMLファイルで実験条件などのパラメータを管理できるようにするPythonライブラリです。

Pythonでyamlの設定ファイルを読み込むようなライブラリとしてはPyYAMLruamel.yamlなどがよく使われます。これらのプリミティブにyamlを読み書きするライブラリと異なり、yacsは実験管理にフォーカスして抽象化されており、主にニューラルネットワークのハイパパラメータ設定などで用いられます。

  • yacsはPyYAMLを使って実装されています
  • ハイパパラメータには、バッチサイズや学習率といった学習時のパラメータに加えて、畳み込み層の入力サイズや深さなどのネットワークのパラメータも含みます

github.com

yacsを使った設定

yacsを使ってパラメータの設定を行っていきます。大まかな流れは4段階です。

  1. デフォルトパラメータのロード
  2. 実験ごとのパラメータで上書き
  3. パラメータを変更できないように固める
  4. パラメータにアクセス

詳細は後述しますが、こういった使い方になります。

from config import get_cfg_defaults

# デフォルトの設定をロード
cfg = get_cfg_defaults()
# 実験ごとのパラメータで上書き
cfg.merge_from_file("experiments/EXP_A.yaml")
# パラメータを変更できないように固める
cfg.freeze()
# パラメータにアクセス
print(cfg.OUTPUT_DIR)

ファイル構成はこんな感じです。

お約束のpip install yacsから始めます。

$ python --version
Python 3.7.4

$ pip install yacs

$ pip list | grep yacs
yacs       0.1.6 

デフォルトパラメータのロード

最初にデフォルト値を生成する処理をPythonで記述します。ファイル名にはconfig.pydefaults.pyが推奨されています (以下ではconfig.pyに実装してます) 。

キーとなるのはCfgNodeで、ここに設定値を持たせることができます。
またCfgNodeは入れ子にすることもできますので、階層構造 (MODELINPUT_SIZE...など) を表現できます。

最後にデフォルト値を設定し終わったCfgNodeのルート (_C) をcloneメソッドで返す関数 get_cfg_defaults を用意しています。

from yacs.config import CfgNode as CN

_C = CN()

_C.OUTPUT_DIR = "output"
_C.LOG_DIR = "log"

_C.MODEL = CN()
_C.MODEL.NAME = "CNN"
_C.MODEL.INPUT_SIZE = [256, 256]

_C.TRAIN = CN()
_C.TRAIN.DEVICE = "cuda:0"
_C.TRAIN.BATCH_SIZE = 16
_C.TRAIN.LR = 0.0001

_C.TEST = CN()
_C.TEST.DEVICE = "cpu"
_C.TEST.BATCH_SIZE = 1

def get_cfg_defaults():
  return _C.clone()

このデフォルト設定を読み込む実装 (main.py) が以下です。

from config import get_cfg_defaults

cfg = get_cfg_defaults()

print(type(cfg))  # <class 'yacs.config.CfgNode'>
print(cfg)
# LOG_DIR: log
# MODEL:
#   INPUT_SIZE: [256, 256]
#   NAME: CNN
# OUTPUT_DIR: output
# TEST:
#   BATCH_SIZE: 1
#   DEVICE: cpu
# TRAIN:
#   BATCH_SIZE: 16
#   DEVICE: cuda:0
#   LR: 0.0001

### 実験ごとにチューニングされたパラメータで上書き 次にこのデフォルト設定を、実験ごとのパラメータで上書きします。この上書きパラメータはYAMLで定義します。ここでは experiments/EXP_A.yaml というファイルに以下のパラメータが記述されているものとします。

TEST:
  DEVICE: "cuda:0"
  BATCH_SIZE: 16

ロードしたCfgNodeオブジェクトのmerge_from_fileメソッドをコールすることで、引数に渡した設定ファイルで自身のパラメータが上書きされます。

  • TESTのBATCH_SIZEとDEVICEが変更されていることに注目です
cfg.merge_from_file("experiments/EXP_A.yaml")

print(cfg)
# LOG_DIR: log
# MODEL:
#   INPUT_SIZE: [256, 256]
#   NAME: CNN
# OUTPUT_DIR: output
# TEST:
#   BATCH_SIZE: 16
#   DEVICE: cuda:0
# TRAIN:
#   BATCH_SIZE: 16
#   DEVICE: cuda:0
#   LR: 0.0001

ちなみにデフォルトに無いキーのパラメータを含む設定ファイルをマージしようとするとエラーになります。想定外のパラメータをアプリケーションでロードできないように保護しています。

NON_EXISTS: "hoge"
cfg.merge_from_file("experiments/EXP_B.yaml")
# -> KeyError: 'Non-existent config key: NON_EXISTS'

パラメータを変更できないように固める

最後にパラメータを変更できないようにfreezeメソッドでイミュータブルにします。freeze後にパラメータの値を変更しようとするとエラーになります。これによって設定ファイル以外からのパラメータの変更を阻止しています。

cfg.freeze()

cfg.LOG_DIR = "hoge"
# -> AttributeError: Attempted to set LOG_DIR to hoge, but CfgNode is immutable

パラメータにアクセス

ロードしたパラメータにアクセスする際は、アトリビュート形式でアクセスします。

print(cfg.OUTPUT_DIR)  # output
print(cfg.MODEL.INPUT_SIZE)  # [256, 256]

まとめ

今回はディープラーニングプロジェクトのパラメータ管理でよく用いられるyacsについて紹介しました。

実際の利用にあたっては デフォルトパラメータのロード -> 実験ごとのパラメータ (YAML) で上書き -> 固定 の手順を間違えやすそうですので、もう1層ラップした関数があると良いかと思います。

def load_config(config_path):
    cfg = get_cfg_defaults()
    cfg.merge_from_file(config_path)
    cfg.freeze()
    return cfg

cfg = load_config("experiments/EXP_A.yaml")