⚙️

OmegaConfで型チェック付きのコンフィグ管理

2021/05/02に公開

機械学習関連でモデルの学習や評価コードを回していると、実行時の設定を管理・保存する必要がある。FacebookResearchからOSSとして公開されているHydraはpythonのコンフィグ管理を行う代表的なツールだが、スクリプト実行時にディレクトリを移動する（YAML上に設定を書き加えることで一応対処可能）などいくつか肌に合わない点があった。

ここでは、特に以下ができるコンフィグ管理を目指す。

前処理、モデル構築、学習などそれぞれのパラメータをまとめて管理できる
パラメータの型チェック
実行時にコマンドラインから与えたパラメータと、デフォルト値と合わせた全体のパラメータを保存できる

最終的なサンプルはこちら

OmegaConfによるコンフィグ構築

OmegaConfはHydraの内部で使われているコンフィグツール。YAMLファイルから設定を読み込んだり、コマンドライン引数をコンフィグオブジェクトに変換したりなど、Hydraのコアになる処理をサポートしている。結果的に、最初に述べたやりたいことはOmegaConfを使って実現できる。

https://github.com/omry/omegaconf

処理ごとのコンフィグと型の付与

前処理、モデル構築などの処理ごとにパラメータをまとめたコンフィグクラスを作る。dataclassを用いて型付きのコンフィグクラスを作ることで、OmegaConfに渡すことができる。。dataclassはPython3.7から追加された機能だが、pip installすることで3.6系でも利用できる。

from typing import Tuple
from omegaconf import OmegaConf
from dataclasses import dataclass


@dataclass
class PreprocessConfig:
    target_size: Tuple[int, int] = (64, 64)
    normalize: bool = True


@dataclass
class ModelConfig:
    image_size: Tuple[int, int] = (64, 64)
    drop_rate: float = 0.5


@dataclass
class TrainConfig:
    epochs: int = 5
    batch_size: int = 32
    learning_rate: float = 0.001


@dataclass
class ExperimentConfig:
    preprocess: PreprocessConfig = PreprocessConfig()
    model: ModelConfig = ModelConfig()
    train: TrainConfig = TrainConfig()
    logdir: str = 'outputs'

dataclass内の各変数に型をつけると、mypyなどの型チェックは機能する（エディタ上でアラートを表示してくれる）一方で、処理自体は実行できてしまう。OmegaConfのAPIを通じてこれらのコンフィグクラスを読み込むことで、各クラス内の変数の型をチェックしてコンフィグを構成してくれる。

base_config = OmegaConf.structured(ExperimentConfig)

とはいえ、上記の例だとデフォルト値を書き間違えない限り型チェックの恩恵はあまりない。個人的に、コマンドラインから引数を設定する時に型チェックが働くと安心できると思う。

コマンドライン引数との結合

OmegaConfを用いることで、コマンドライン引数の取得と前述のデフォルトコンフィグを以下のように簡単に結合できる。

base_config = OmegaConf.structured(ExperimentConfig)
cli_config = OmegaConf.from_cli()
config = OmegaConf.merge(base_config, cli_config)

コマンドライン引数はHydraと同様に、$ python train.py train.epochs=10などといった形で渡すことができる。この時に例えばtrain.epochs=0.5などとすると、定義した型と異なるためエラーを吐いてくれる。

$ python train.py train.epochs=0.5
(--中略--)
omegaconf.errors.ValidationError: Value '0.5' could not be converted to Integer
        full_key: train.epochs
        reference_type=TrainConfig
        object_type=TrainConfig

間違えて存在しないパラメータ（epochsをepochなど）を指定してしまってもエラーを出してくれる。

$ python train.py train.epoch=10
(--中略--)
omegaconf.errors.ConfigKeyError: Key 'epoch' not in 'TrainConfig'
        full_key: train.epoch
        reference_type=TrainConfig
        object_type=TrainConfig

コンフィグの保存

コンフィグの保存は以下のように可能。Hydraと同様にコマンドラインから上書きした値と、結合後の全体のコンフィグを両方保存することで、各実行の差分がわかりやすくなる。

OmegaConf.save(config, 'config.yaml')
OmegaConf.save(cli_config, 'override.yaml')

これで一通りやりたいことができた。全体のコードはこちら

その他

値の補完

OmegaConfの機能として、あるコンフィグで用いた変数の値を別のコンフィグに用いることができる。例えば上の例のModelConfig.image_sizeをPreprocessConfig.target_sizeと同じにしたい場合には、以下のように実装できる。

from typing import Tuple
from omegaconf import OmegaConf
from dataclasses import dataclass


@dataclass
class PreprocessConfig:
    target_size: Tuple[int, int] = (64, 64)
    normalize: bool = True


@dataclass
class ModelConfig:
    image_size: Tuple[int, int] = '${preprocess.target_size}'
    drop_rate: float = 0.5

preprocess.target_sizeによる参照は各コンフィグを結合する（今回の例だとOmegaConf.structured(ExperimentConfig)）ことで可能になる。

また上記のまま静的型チェックを行うとTuple[int, int] != stringとして判定されてしまうが、これはomegaconf.IIを被せることで回避できる（処理としてはstring型の引数をAny型として返しているだけ）。

from omegaconf import OmegaConf, II

@dataclass
class ModelConfig:
    image_size: Tuple[int, int] = II('preprocess.target_size')
    drop_rate: float = 0.5

ただし、この状態で実行時にmodel.image_size=\[32,32\]と指定すると、preprocess.target_sizeは変わらないままmodel.image_sizeだけ上書きされるので注意が必要。そもそも確実に同じ値を用いるならプログラム内で処理する方がよさそう。

可変なデフォルト値

dataclassの変数の型を例えばリストで定義しようとすると次のエラーが出る。

@dataclass
class PoorClass:
    x: List[int] = [1, 2, 3]

ValueError: mutable default <class 'list'> for field x is not allowed: use default_factory

これは公式ドキュメントにも記載されている。list, dict, setをフィールド変数に用いたい場合は、言われている通りdefault_factoryを使って以下のように書けば良い。

from dataclasses import dataclass, field

@dataclass
class GoodClass:
    x: List[int] = field(default_factory=lambda: [1,2,3])

ちなみに（これもドキュメントに書いてあるが）、dataclassを使わずにlist型のフィールド変数を持つクラスを作ると次のようなことが起こる。

class PoorClass:
    x: List[int] = [1, 2, 3]
    
a = PoorClass()
b = PoorClass()
a.x is b.x  # => True

a.x, b.xがどちらも同じオブジェクトを参照していることがわかる。これはPythonの仕様だが、多くの場合であまり直感的ではない。dataclassではこの問題を回避するためにデフォルトファクトリ関数の使用を促すエラーを出力している。

感想

冒頭に挙げた項目をカバーして、程よく型の恩恵を受けられている気がする。一方でこのままだと-h, --helpでヘルプを表示できないのが地味に悲しい。また今回のサンプルだと前処理、モデル、学習の3つのコンフィグクラスを作っているだけだが、階層が増えてくるとその分たくさんコンフィグクラスを作る必要がある点も課題。

OmegaConfはYAMLファイルも読み込むことができるので、これと組み合わせた使い方もありかもしれない。型の実装・指定自体はPython側でカバーした方がいいと思う。コンフィグの構成（前処理、モデル、学習など）があまり変更されない場合はそれぞれコンフィグクラスを実装して、構成自体を頻繁にいじる場合（Kaggleとか）ではやはりYAMLで管理するのが手軽に思う。