設定管理ツール Hydra で内部構造ごと書き換える。

はじめに

本記事は 機械学習の領域で主にハイパーパラメータの管理に用いられるツールである、Hydraについて解説する。Hydraは単にハイパーパラメータを階層的に記述できるというだけでなく、より複雑に、クラス・処理構造そのものを、ソースファイルを新たに追加することなく変更したい、という要望に応える。

本記事の対象読者はある程度Pythonやシステム設計に精通した者とする。モジュールシステム、クラス、デコレータなどの概念を理解していることが前提である。

さらに本記事は 機械学習のハイパーパラメータの管理だけでなく、システムそのものを コンフィギュレーション することに興味がある読者も対象となる。本記事の知見は機械学習以外のCS領域における実験系でも活用できるだろう。

ただし、シンプルな用途においてはオーバーエンジニアリングになる恐れがあることに注意してほしい。本記事で紹介する方法はあくまでも複雑なシステムのコンフィギュレーションを行うためのものであることを意識してほしい。

また、本記事を理解する上で必要最低限度のコードスニペットは用意したが、より適切な理解を得たいのであれば、lightning-hydra-templateを参照すると良い。筆者もこのリポジトリを参考にしている。

前置きは以上である。それでは、始めよう。

Hydra とは

基本的な使用方法に関しては以下に示す公式ドキュメントを参照願う。

https://hydra.cc/docs/intro/

Hydraは、研究やその他の複雑なアプリケーションの開発を簡素化する、オープンソースのPythonフレームワークです。主な特徴は、構成ファイルやコマンドラインを通じて、階層的な構成を動的に作成し、それを合成して上書きする能力にあります。Hydraという名前は、複数の類似したジョブを実行できるその能力から来ており、まるで多頭のヒドラのようです。

主な特徴:

• 複数のソースから構成可能な階層的構成
• コマンドラインから指定または上書き可能な構成
• 動的なコマンドラインタブ補完
• アプリケーションをローカルで実行するか、リモートで実行するために起動
• 単一のコマンドで異なる引数を持つ複数のジョブを実行

つまり、ちょっと学習率 lrを変えて実験してみたいときに、ファイルを編集せずに コマンドラインから実行時に変更できたり、大きく設定を変更してGitHubにpushしたいが既存の設定ファイルは残しておきたいときに 既存設定を引き継ぎつつ新しく定義する ことができると言っている。さらに複数パターン試したいときは全部自動で実行してくれる。

Trial & Error のサイクルを高速化し、より効率的な実験活動を促進するツールである。便利すぎる (さすが Meta社)

エントリーポイント

実行するエントリーポイントとなる関数に対して、@hydra.mainデコレータをつける。その実行ファイルの位置を起点に、conf ディレクトリ内部の、config.yamlが自動的に読み込まれ、omegaconfのDictConfigオブジェクトが渡される。

conf/config.yaml

db:
  driver: mysql
  user: omry
  pass: secret

my_app.py

import hydra
from omegaconf import DictConfig, OmegaConf

@hydra.main(version_base=None, config_path="conf", config_name="config")
def my_app(cfg : DictConfig) -> None:
    print(OmegaConf.to_yaml(cfg))

if __name__ == "__main__":
    my_app()
# db:
#   driver: mysql
#   pass: secret
#   user: omry

コマンドラインから設定を上書きすることもできる。とても便利。

$ python my_app.py db.user=root db.pass=1234
db:
  driver: mysql
  user: root
  pass: 1234

基本的な使い方は公式ドキュメントに書いてあるのでそこを見たほうがよい。
https://hydra.cc/docs/intro/#basic-example

オブジェクトのインスタンス化

https://hydra.cc/docs/advanced/instantiate_objects/overview/

本記事で最初に取り上げる機能は、設定ファイルから行うオブジェクトの インスタンス化である。機械学習との関連性をわかりやすくするために、PyTorchを用例に用いる。

まずは、設定ファイルから線形層 (Linear)を呼び出してみよう。
_target_特殊キーにそのオブジェクトまでの python dot pathを指定し、引数名とその値をyamlに記述する。（これはキーワード引数を指定している）

instantiation.py

import hydra
import omegaconf


yaml1 = """
_target_: torch.nn.Linear
in_features: 2
out_features: 3
"""

conf1 = omegaconf.OmegaConf.create(yaml1)
linear1 = hydra.utils.instantiate(conf1)
print(type(linear1), linear1)

output

<class 'torch.nn.modules.linear.Linear'> Linear(in_features=2, out_features=3, bias=True)

---

_args_ 特殊キーを使うと ポジショナル引数として渡せる。 nn.Sequentialなどを使うときに便利。

yaml2 = """
_target_: torch.nn.Sequential
_args_:
  - _target_: torch.nn.Linear
    _args_: [16, 32]
  - _target_: torch.nn.Linear
    _args_: [32, 8]
"""

conf2 = omegaconf.OmegaConf.create(yaml2)
seq2 = hydra.utils.instantiate(conf2)
print(seq2)

output

Sequential(
  (0): Linear(in_features=16, out_features=32, bias=True)
  (1): Linear(in_features=32, out_features=8, bias=True)
)

---

_partial_: True とするとfunctool.partialでラップされ、インスタンス化するときに渡す引数を部分的に指定した状態で返すことができる。ほかのオブジェクトとの依存関係が存在する場合に便利である。（依存関係を解消する処理は pythonソース上で作る必要があることに注意）
またあえてすべての引数を指定してpartialにすれば、引数なしで何度でもインスタンスを生成可能なオブジェクトを作ることができる。

yaml3 = """
optim: 
  _target_: torch.optim.Adam
  _partial_: True
  lr: 0.001
  betas: [0.9, 0.999]

model:
  _target_: torch.nn.Linear
  _args_: [10, 20]
"""

conf3 = omegaconf.OmegaConf.create(yaml3)
optim3_partial = hydra.utils.instantiate(conf3.optim)
model3 = hydra.utils.instantiate(conf3.model)

print("Partial:", optim3_partial, "\n")
optim3 = optim3_partial(model3.parameters())
print("Complete Instance:",optim3)

output

Partial: functools.partial(<class 'torch.optim.adam.Adam'>, lr=0.001, betas=[0.9, 0.999])

Complete Instance: Adam (
Parameter Group 0
    amsgrad: False
    betas: [0.9, 0.999]
    capturable: False
    differentiable: False
    eps: 1e-08
    foreach: None
    fused: None
    lr: 0.001
    maximize: False
    weight_decay: 0
)

ちなみに _partial_ を用いずとも、instantiateを呼び出す際に引数としてインスタンス化に必要なオブジェクトがすべて渡されていれば、問題なくインスタンス化できる。

yaml4 = """
optim: 
  _target_: torch.optim.Adam
  lr: 0.001
  betas: [0.9, 0.999]

model:
  _target_: torch.nn.Linear
  _args_: [10, 20]
"""

conf4 = omegaconf.OmegaConf.create(yaml4)
model4 = hydra.utils.instantiate(conf4.model)
optim4 = hydra.utils.instantiate(conf4.optim, model4.parameters())
print("Complete Instance:",optim4)

---

他の特殊キーとしては、再帰的なインスタンス化を制御する_recursive_フラグや、_convert_オプションなどがある。

さらにクラスオブジェクトやモジュール内に定義された定数を取得するための便利なメソッドも用意されており、_target_にhydra.utils.get_classやhydra.utils.get_objectを指定し、その引数であるpathにそのオブジェクトまでの dot pathを指定することで取得することができる。

tensor_type:
  _target_: hydra.utils.get_class
  path: torch.Tensor
dtype:
  _target_: hydra.utils.get_object
  path: torch.bfloat16

詳しくは hydraのリファレンスを参照願う。

https://hydra.cc/docs/advanced/instantiate_objects/overview/

---

こういったことができるといったい何がうれしいのだろうか?

それは、pythonソースファイルを編集せずに、instantiateするオブジェクトを変更して処理構造を書き換えられることだ。

複雑な機械知能系において、その内部構造をとっかえひっかえしたくなることは、かなりある。 そういった際に、pythonソースファイルでなくyaml設定ファイルから記述することによって、大量に実験用ソースファイルが散乱することを避けることができる（設定ファイルは散乱するが）

ソースファイルは再利用性高く書き、何度も書き換える部分を設定ファイルに任せることにより、継続的な開発とラピッドな実験を可能にする。素晴らしい。

設定ファイルの階層的管理

Hydraは階層的に設定ファイルを配置し、その構造を保持した状態で読み込むことができる。これにより、一つの巨大な設定ファイルを小さな設定ファイルに分割することが可能になる。再利用性や設定ファイルの可読性を向上させることができる。

次に階層的な設定ファイルの例を示す。

ディレクトリ構造

.
├── configs
│   ├── data
│   │   └── mnist.yaml
│   ├── model
│   │   └── conv.yaml
│   └── train.yaml
└── train.py

configs/data/mnist.yaml

data_dir: path/to/mnist_data_dir
batch_size: 128

configs/model/conv.yaml

_target_: torch.nn.Conv2d
in_channels: 1
out_channels: 8
kernel_size: 3

configs/train.yaml

defaults:
  - _self_
  - data: mnist
  - model: conv

task_name: train_on_mnist

train.py

import hydra
from omegaconf import DictConfig, OmegaConf

@hydra.main(version_base="1.3", config_path="./configs", config_name="train")
def main(cfg: DictConfig) -> None:
    print(OmegaConf.to_yaml(cfg))


if __name__ == "__main__":
    main()

---

指定した階層にある設定ファイルが読み込まれ、次のように一つの設定オブジェクトが構築される。
defaultsに指定した順序で設定ファイルが読み込まれる。同じ階層・名前の設定値が存在した場合は後から読み込まれた値で上書きされる。_self_はそれが記述された設定ファイル自体を指す。今回の例ではtask_nameが一番最初に読み込まれ、続いてdata, modelが読み込まれる。

output

$ python train.py
task_name: train_on_mnist
data:
  data_dir: path/to/mnist_data_dir
  batch_size: 128
model:
  _target_: torch.nn.Conv2d
  in_channels: 1
  out_channels: 8
  kernel_size: 3

他の設定ファイルで値を上書きする。

他の設定ファイルから値を上書きする目的は、基礎となる設定項目を継承しつつ、特定の項目を変化させ実験するためだ。

ここで、configs/experiment/example.yamlを作成し、次のように記述する。

configs/experiment/example.yaml

# @package _global_

task_name: "example_override"

python train.py に、+experiment=exampleと追加すると、task_nameが上書きされる。

$ python train.py +experiment=example
task_name: example_override
data:
  ... # 省略
model:
  ... # 省略

ここで、example.yamlの冒頭に、# @package _global_ と追加されていることに注意してほしい。これは設定ファイルによって指定される値の名前空間を指定するディレクティブである。通常はexperiment名前空間に所属するexample.yamlは、_global_ ディレクティブによって、グローバル名前空間、すなわちルートから設定するようになる。

https://hydra.cc/docs/advanced/overriding_packages/

---

train.yamlに明示的にexperiment=nullで指定しておけば、上書きする優先順位や、実行時に+をつけなくても良くなる。(i.e. python train.py experiment=example)

train.yaml

defaults:
  - _self_
  - data: mnist
  - model: conv
  - experiment: null

task_name: train_on_mnist

もちろん、ある特定の階層の設定をまるごと変更することも可能だ。override <階層>: <ターゲット設定名>と指定する。（値の上書きの関係上 デフォルトリストと組み合わせて使っている。）

configs/experiment/example.yaml

# @package _global_

defaults:
  - override /data: cifar10

task_name: "example_override"

別の場所の設定値を参照し、補完する。

hydraというよりも omegaconfの機能であるが、設定内の別の箇所の値を参照し、同じ値を用いることができる。その参照先は${dot.path}によって指定する。

interpolation

name1: 10
name2:
  name2_1: ${..name1} # 10

ドット.をつけるとその設定階層からの相対参照となり、ドットの数 $N - 1$ だけ遡って参照する。

ドット.をつけなければ構造上のルートから辿って値を参照する。グローバルに参照したい値を特定の設定ファイルにまとめておき、そこから参照する際に便利である。

また環境変数や独自のResolverを用いて補完することもできる。hydraによっていくつか提供されている。ocはomegaconfから提供されているものを指す。

root_dir: ${oc.env:PROJECT_ROOT}
output_dir: ${hydra:runtime.output_dir}
work_dir: ${hydra:runtime.cwd}

https://hydra.cc/docs/1.3/configure_hydra/intro/#resolvers-provided-by-hydra
https://omegaconf.readthedocs.io/en/latest/custom_resolvers.html

ログをカスタマイズする。

hydraはloggingモジュールの設定も行うことができる。知っておくとそこそこ便利である。
hydra/job_loggingに loggingモジュールのDictConfigに設定可能なyamlを記述する。

hydra:
  job_logging:
    handlers:
      file:
        filename: ${hydra.runtime.output_dir}/train.log

https://hydra.cc/docs/configure_hydra/logging/

まとめ

本記事では いくつかの特筆すべきHydraの機能を述べた。エントリーポイントによりシステム起動時にどのように設定ファイルが読み込まれるかを記述し、システムそのものをコンフィギュレーションすることが可能となるオブジェクトのインスタンス化機能について述べた後、値の参照補完やloggingモジュールの設定といったシステム構築に便利な機能を紹介した。

Hydraは機械学習のハイパーパラメータ管理としてももちろん有用であるが、より複雑かつ実験的な機械知能システムを実装する上で、非常に強力なツールである。

本記事でその一端が伝われば幸いである。

参考

• hydra: https://hydra.cc/
• omegaconf: https://omegaconf.readthedocs.io/en/latest/index.html
• lightning-hydra-template: https://github.com/ashleve/lightning-hydra-template