argparse-dataclassからhydraに移行する

argparse-dataclassを少し使っていたのだが、hydraに移行してみたのでメモ。

前回書いたこちらの話に関係するが、解決したもののやはりhydraを試してみるかと思ったため。
https://jsapachehtml.hatenablog.com/entry/2021/10/23/184058

題材としては今年始めくらいにやっていたテキスト感情分類のrepoを使う。hydraを導入した際の差分はこちら。
https://github.com/y-kamiya/emotion-classification/pull/11

元々argparse-dataclassを使っていたためdataclassでオプションのスキーマが定義されている状態からスタートしているのもありStructured Configsの形で入れた。導入に必要な情報はこちらのStructured Configs用のチュートリアルを見ればOK。
https://hydra.cc/docs/tutorials/structured_config/intro

この形にしておくことで、dataclassとしてオプションのスキーマを定義することができ、それによって型による入力値のチェックが可能になる。あとそもそもオプションがクラスの形で定義されている方が実装上便利。

hydra全体の使い方については既にわかりやすく書かれたブログなどたくさん上がっているのでそちらを参考にするとして、ここでは移行する際に調べたり考えたりした点を書いておく。

argparseにおけるchoices

hydraではargparseにおけるchoicesの指定そのままのものはない(特定の文字列のみを引数として許す機能)。ただし、スキーマの型指定が可能であるためenumを利用することで実装的にはより使いやすく定義できる。log levelを取れるようにしていたのでそれに関する部分だけ例として書くと

# argparse-dataclassの場合
class Config:
    loglevel: str = field(
        default="INFO",
        metadata=dict(choices=["DEBUG", "INFO", "WARN", "ERROR", "CRITICAL"]),
    )
# hydraの場合
class LogLevel(Enum):
    DEBUG = auto()
    INFO = auto()
    WARN = auto()
    ERROR = auto()
    CRITICAL = auto()

class Config:
    loglevel: LogLevel = LogLevel.INFO

hydraでは上記の書き方にしておくことで実行の際は以下のように指定できる。

$ python main.py loglevel=DEBUG

元々choicesを使ったオプションの指定がいくつかあり、その判定のためにオプションの値となる文字列との比較がコード内に存在していた。いずれenumにしようかと思っていたが今回の対応のおかげで必然的に修正されることに。

Noneの扱い

hydraのStructured Configsでは例えば str と指定したオプションはデフォルト値としての指定も含めてNoneを入れることができない。

Noneを取れるようにするには以下のようにする必要がある。

from typing import Optional

class Config:
    example: Optional[str] = None

ただ、元々のオプションの中でNoneを使っていたものを見てみるとそのほとんどが良くない使い方をしてる部分だった。本来的にはconfigに入れるべきではないオブジェクトを取り回ししやすいようにconfigに入れてただけ。こちらのloggerとか。
https://github.com/y-kamiya/emotion-classification/pull/11/files#diff-e9f42701dc7658ddc2df3f56acc27573667c95b86723dcbe10ceb77ece3ac793L54

私は基本的にtrainerやdataset用のクラスなどすべてにconfigオブジェクトを渡してアクセス可能な形で実装するようにしていたため、configに入れておくと手間がかからないというだけだった。これもいずれ、、と思いつつ案の定そのままになっていたものなのでこれを機にconfigから外した。

また、hydraの制約としてもそのような使い方はできないようになっており、Structured Configsとして正しく解釈されるのは以下の型のみで他はエラーとなる。

  • Primitive types (int, bool, float, str, Enums)
  • Nesting of Structured Configs
  • Containers (List and Dict) containing primitives or Structured Configs
  • Optional fields

cf. https://hydra.cc/docs/tutorials/structured_config/intro#structured-configs-supports

post_initについて

dataclassでは __post_init__ というメソッドを定義することで __init__の後にすべき処理を定義できる。
https://docs.python.org/ja/3/library/dataclasses.html#post-init-processing

hydraではconfigが生成される際にこちらの処理は反映されていなかった。少し上で書いたカスタムクラスのオブジェクトを保持する際に、その初期化に用いていたので実質使う必要がなくなったため問題ないが一応メモしておく。

別の設定値を参照

configに存在する値を参照して別の値を設定することができる。
https://hydra.cc/docs/patterns/specializing_config#modified-configyaml

defaults:
  - dataset: imagenet
  - model: alexnet
  - optional dataset_model: ${dataset}_${model}

これはOmegaConfのvariable interpolationという機能をそのまま使ったものなので詳細はこちら。
https://omegaconf.readthedocs.io/en/latest/usage.html#variable-interpolation

注意点としては文字列として挿入されることになるということ。boolとして扱おうとしてうまくいかなかった例
https://github.com/y-kamiya/emotion-classification/commit/e83b8e863fd6bb1c85c0aff89f2c76cea8dae6e2#diff-e9f42701dc7658ddc2df3f56acc27573667c95b86723dcbe10ceb77ece3ac793L57

自動でcwdが変更される

hydraは出力を実行毎に分けるために実行時に current working directory が変更されるようになっている。
https://hydra.cc/docs/tutorials/basic/running_your_app/working_directory

上記のページにある例だが以下のように変化する。

  • 実行前
    • /home/omry/dev/hydra
  • 実行中
    • /home/omry/dev/hydra/outputs/2019-09-25/15-16-17

入力データなどをpathで渡して使っていたため、この挙動によってデータが取得できなくなるという問題が起きた。

出力されるデータは毎回上書きしてしまってOKなのであれば設定でディレクトリがそのままになるようにすることが可能。
https://hydra.cc/docs/configure_hydra/workdir/

hydra:
  run:
    dir: .

output用のディレクトリ構成は使いたいという場合は上記の設定は変更せず、コード上で元のcwdを参照するように変更する。以下のメソッドで元のを取得可能。
https://hydra.cc/docs/tutorials/basic/running_your_app/working_directory#original-working-directory

ちなみに私は基本的にconfig全体を各クラスに引き渡して参照する形で実装していたので、mainの最初の部分でconfig内の値を修正するようにした。
https://github.com/y-kamiya/emotion-classification/blob/6560332/src/main.py#L22-L24

default list

同じ項目の設定が複数の場所に存在する場合にどれを優先するかの設定。
https://hydra.cc/docs/advanced/defaults_list

こちらに2つの設定の順番を変えた場合にどちらが優先されるかが書いてあるが、後ろに設定したもので上書きするというイメージ。
https://hydra.cc/docs/advanced/defaults_list#composition-order

設定はyamlならファイルのパス、スキーマであればCacheStoreの登録時に設定している名前を使う。
https://github.com/y-kamiya/emotion-classification/blob/6560332/src/main.py#L17

yamlで設定している内容を優先して使いたいため、このようなdefault listを定義している。
https://github.com/y-kamiya/emotion-classification/blob/6560332/src/conf/main.yaml#L1-L3

tab completion

シェル上で引数指定する際にtab補完を効かせることができる。
https://hydra.cc/docs/tutorials/basic/running_your_app/tab_completion

これはできると便利なのでやってみたのだが、tabを押すとしばらく動かなくなるくらい重くなってしまったので一旦やめた。

まとめ

今回書いていない機能も当然あるため、argparseに比べるとかなり高機能でありその分複雑であるといえる。今回のように一度やってしまえば特にわかりづらい部分はないので問題ないが初見だと追いづらい部分もある(どの設定が優先されるのかなど)

また、基本的にはOmegaConfの機能を利用してるので何かやったり引っかかったりした場合はOmegaConfの使い方を調べるのがよい。
- https://omegaconf.readthedocs.io/en/2.1_branch/index.html

参考