はじめに
今回は Hydra を使った設定管理について調べてみました！
Hydra というのは Meta が開発している python のフレームワークの一つであり、主に設定ファイルの管理に長けているものです。
この記事では、いくつかの基本的はシチュエーションにおける Hydra をみていきたいと思っています 🐲 🐲 🐲
1. yaml で設定管理したい
こちらの記事にあるように、argparse を使いプログラム実行時の引数を受け取る方法はよく使われていると思います。しかし、設定するパラメータ数が多いときなどには苦しさを感じることも多々あります。
以下の例で、特に設定数は多くないですが hydra を使ってみようと思います! hugguingface trainer を用いた fine-tuning のサンプルコードを参考にして例を作成しています。
import argparse import transformers from datasets import load_dataset from transformers import AutoModelForSequenceClassification, AutoTokenizer def main(): parser = argparse.ArgumentParser(description="huggungface transformers training") # transformers parser.add_argument("--model_name", type=str, default="bert-base-uncased") parser.add_argument("--num_labels", type=int, default=2) parser.add_argument("--per_device_train_batch_size", type=int, default=8) parser.add_argument("--per_device_eval_batch_size", type=int, default=8) parser.add_argument("--evaluation_strategy", type=str, default="epoch") parser.add_argument("--num_epochs", type=int, default=3) parser.add_argument("--learning_rate", type=float, default=5e-5) parser.add_argument("--warmup_ratio", type=float, default=0.1) parser.add_argument("--gradient_accumulation_steps", type=int, default=1) parser.add_argument("--eval_accumulation_steps", type=int, default=1) parser.add_argument("--weight_decay", type=float, default=0.01) parser.add_argument("--save_strategy", type=str, default="epoch") parser.add_argument("--fp16", type=bool, default=False) # paths parser.add_argument("--logging_dir", type=str, default="logs") parser.add_argument("--output_dir", type=str, default="output") args = parser.parse_args() # model, tokenizer のロード tokenizer = AutoTokenizer.from_pretrained(args.model_name) model = AutoModelForSequenceClassification.from_pretrained( args.model_name, num_labels=args.num_labels ) # example データセットのロード raw_datasets = load_dataset("glue", "mrpc") def tokenize_function(example): return tokenizer(example["sentence1"], example["sentence2"], truncation=True) tokenized_datasets = raw_datasets.map(tokenize_function, batched=True) # トレーニングの設定 training_args = transformers.TrainingArguments( output_dir=args.output_dir, per_device_train_batch_size=args.per_device_train_batch_size, per_device_eval_batch_size=args.per_device_eval_batch_size, evaluation_strategy=args.evaluation_strategy, logging_dir=args.logging_dir, num_train_epochs=args.num_epochs, learning_rate=args.learning_rate, warmup_ratio=args.warmup_ratio, gradient_accumulation_steps=args.gradient_accumulation_steps, eval_accumulation_steps=args.eval_accumulation_steps, weight_decay=args.weight_decay, save_strategy=args.save_strategy, fp16=args.fp16, ) trainer = transformers.Trainer( model, training_args, train_dataset=tokenized_datasets["train"], eval_dataset=tokenized_datasets["validation"], tokenizer=tokenizer, ) trainer.train() if __name__ == "__main__": main()
デフォルト値は設定していますが、全引数を設定すると以下のようになります。
rye run python src/main_argparse.py --model_name bert-base-uncased --num_labels 2 --per_device_train_batch_size 8 --per_device_eval_batch_size 8 --evaluation_strategy epoch --num_epochs 3 --learning_rate 5e-5 --warmup_ratio 0.1 --gradient_accumulation_steps 1 --eval_accumulation_steps 1 --weight_decay 0.01 --save_strategy epoch --fp16 False --logging_dir logs --output_dir output
上記のような設定を yaml ファイルに書き、hydraを使うことで argparse から脱却することができます。また、設定ファイル自体に階層構造を持たせる、つまりグループ化することで、よりわかりやすく管理しやすい形で設定ファイルを扱うことができます。
今回は configs ディレクトリを作成し、さらにその中に paths と transformers を作り、それぞれに対応する設定ファイルを作ろうと思います。
├── configs
│ ├── config.yaml
│ ├── paths
│ │ └── default.yaml
│ └── transformers
│ └── default.yaml
└── src
├── __init__.py
└── main_hydra.py
default.yaml には対応するディレクトリ名に関するデフォルト値を記載し、config.yaml は、それぞれの設定ファイルをまとめる役割を持っています。
model_name: bert-base-uncased num_labels: 2 per_device_train_batch_size: 8 per_device_eval_batch_size: 8 evaluation_strategy: epoch num_epochs: 3 learning_rate: 2e-5 warmup_ratio: 0 gradient_accumulation_steps: 1 eval_accumulation_steps: 1 weight_decay: 0.01 save_strategy: epoch fp16: False
logging_dir: logs output_dir: output
defaults: - paths: default - transformers: default - _self_
defaults: に関しては公式ドキュメントを参照ください。ディレクトリ名:設定ファイル名 という形式で設定を読み込み、指定した設定を使用できるようにします。
また config.yaml 自体にも paths や transformers の設定以外の設定などを書くことができます。- _self_ は、config.yaml 自身の設定を明示的に表しているものに過ぎません。
ただ、同じ設定が存在する場合はリストのより後のものが優先されます。今回の場合だと _self_ が最優先ということになります。
実行対象のファイルは以下のようになります。 main() に対して @hydra.main デコレータの追加があります。最終的に使う設定ファイルと、その設定ファイルが存在するディレクトリのパスをここで指定することで、その設定が使えるようになります。
import hydra import transformers from datasets import load_dataset from omegaconf import DictConfig from transformers import AutoModelForSequenceClassification, AutoTokenizer @hydra.main(config_path="../configs", config_name="config", version_base="1.3") def main(cfg: DictConfig): # model, tokenizer のロード tokenizer = AutoTokenizer.from_pretrained(cfg.transformers.model_name) model = AutoModelForSequenceClassification.from_pretrained( cfg.transformers.model_name, num_labels=cfg.transformers.num_labels ) # example データセットのロード raw_datasets = load_dataset("glue", "mrpc") def tokenize_function(example): return tokenizer(example["sentence1"], example["sentence2"], truncation=True) tokenized_datasets = raw_datasets.map(tokenize_function, batched=True) # トレーニングの設定 training_args = transformers.TrainingArguments( output_dir=cfg.paths.output_dir, per_device_train_batch_size=cfg.transformers.per_device_train_batch_size, per_device_eval_batch_size=cfg.transformers.per_device_eval_batch_size, evaluation_strategy=cfg.transformers.evaluation_strategy, logging_dir=cfg.paths.logging_dir, num_train_epochs=cfg.transformers.num_epochs, learning_rate=cfg.transformers.learning_rate, warmup_ratio=cfg.transformers.warmup_ratio, gradient_accumulation_steps=cfg.transformers.gradient_accumulation_steps, eval_accumulation_steps=cfg.transformers.eval_accumulation_steps, weight_decay=cfg.transformers.weight_decay, save_strategy=cfg.transformers.save_strategy, fp16=cfg.transformers.fp16, ) trainer = transformers.Trainer( model, training_args, train_dataset=tokenized_datasets["train"], eval_dataset=tokenized_datasets["validation"], tokenizer=tokenizer, ) trainer.train() if __name__ == "__main__": main()
cfg.transformers.model_name や cfg.paths.output_dir のように ディレクトリ名.パラメタ名 という感じでアクセスできます。
引数をつけずに実行すると、もちろん yaml ファイルに書いた通りの設定で実行されます。
引数設定の例は以下です。
python src/main_hydra.py transformers.fp16=true
また、実行時にはデフォルトで outputs フォルダが作成されます。この中には log ファイルや、実行時のすべての hydra の設定ファイルが実行日/実行時間のフォルダに保存されます。
2. クラス単位で設定したい
設定パラメタを受け取り、あるクラスや関数の引数として使用する場合、なんとなく冗長な気がします。また、パラメタの値によって対象のクラス・関数を変更したい時、いちいち対象のクラス・関数を import し条件文で分岐を作るなども少しダルイ感じがあります。
hydra には Instantiating というシステムがあり、これが前述の問題を解決してくれます。
これを使うことで最終的な main.py は以下のようになります。
import hydra from datasets import load_dataset from omegaconf import DictConfig @hydra.main(config_path="../configs", config_name="config", version_base="1.3") def main(cfg: DictConfig): # model, tokenizer のロード model = hydra.utils.get_method(cfg.transformers.model)( cfg.transformers.model_name, num_labels=cfg.transformers.num_labels, ) tokenizer = hydra.utils.get_method(cfg.transformers.tokenizer)(cfg.transformers.model_name) # example データセットのロード raw_datasets = load_dataset("glue", "mrpc") def tokenize_function(example): return tokenizer(example["sentence1"], example["sentence2"], truncation=True) tokenized_datasets = raw_datasets.map(tokenize_function, batched=True) trainer = hydra.utils.instantiate( cfg.transformers.trainer, model=model, train_dataset=tokenized_datasets["train"], eval_dataset=tokenized_datasets["validation"], tokenizer=tokenizer, ) trainer.train() if __name__ == "__main__": main()
注目すべきは trainer = hydra.utils.instantiate( こちらですね。configs/transformers/default.yaml にある trainer を instantiate しています。
さらには model や tokenizer さえも yaml で設定できます。今回は AutoTokenizer などの Auto 系のクラスがあるのでいいですが、特定のクラスを使うときなどには get_method や get_class が使えます。
configs/transformers/default.yaml はこんな感じです。
model_name: bert-base-uncased num_labels: 2 model: transformers.AutoModelForSequenceClassification.from_pretrained tokenizer: transformers.AutoTokenizer.from_pretrained trainer: _target_: transformers.Trainer args: _target_: transformers.TrainingArguments output_dir: ${paths.output_dir} logging_dir: ${paths.logging_dir} per_device_train_batch_size: 8 per_device_eval_batch_size: 8 evaluation_strategy: epoch num_train_epochs: 3 learning_rate: 2e-5 warmup_ratio: 0 gradient_accumulation_steps: 1 eval_accumulation_steps: 1 weight_decay: 0.01 save_strategy: epoch fp16: False
_target_ に instantiate 対象へのパスを書き、その下に引数を書くことができます。ここにすべての引数を書かずとも前述のように
trainer = hydra.utils.instantiate(
cfg.transformers.trainer,
model=model,
train_dataset=tokenized_datasets["train"],
eval_dataset=tokenized_datasets["validation"],
tokenizer=tokenizer,
)
インスタンス化のタイミングで追加で引数を設定できます。また、functools.partial と同様のものも作成することができます。(参考)
hydra.utils.get_class や hydra.utils.get_method を使うことで、クラスやそのメソッド、関数そのものを呼び出すことができます。
3. 設定をカスタマイズしたい
bert 系のモデルを使う時など、large モデルを使うときはマシンの関係上バッチサイズも小さく、さらに学習率も小さくしたいなど、対象のモデルに応じて各設定のデフォルトを変更したい場合があります。もちろんコマンドライン上で transformers.model_name=roberta-large transformers.trainer.args.learning_rate=1e-5 などのように override することもできますが、設定ファイルとして残したい気持ちもあります。
├── configs
│ ├── config.yaml
│ ├── paths
│ │ └── default.yaml
│ └── transformers
│ |── default.yaml
| |── roberta-base.yaml
| └── roberta-large.yaml
└── src
├── __init__.py
└── main.py
このように roberta-base.yaml と roberta-large.yaml の設定ファイルを追加してみます。
defaults: - default model_name: roberta-base
defaults: - default model_name: roberta-large trainer: args: num_train_epochs: 2 learning_rate: 1e-5 gradient_accumulation_steps: 4
defaults: - default は、同じ階層の default.yaml を overriede するために追加します。それぞれの設定ファイルからみて取れるように、変更点 (override 対象) 以外は、 defaults で指定した値のパラメタが適用されます。この場合だと defaults.yaml がそれですね。
こうすることで、変更箇所以外の余計な設定を省略しつつ設定を行うことができました。
次にこの設定を適用する方法についてです。
configs/config.yamlを書き換える- 実行時のコマンドライン引数で指定する
1 については、以下のように config.yaml を修正します。default だった部分を roberta-large にしただけです。
defaults: - paths: default - transformers: roberta-large - _self_
2 については以下のコマンドライン引数を使います。transformers の設定を変更した感じです。
python run src/main.py transformers=roberta-large
話はずれますが、hydra ではデフォルトの値を設定せず、コマンドライン引数として必ず設定するパラメタは ??? と書くことができます。??? に当たる部分が未指定だとエラーが発生します。
defaults: - paths: default - transformers: ??? - _self_
例えば上のように書けば、python run src/main.py transformers= をコマンドラインから指定する必要があります (コマンドライン以外からも指定する方法はあります)。
4. その他
環境変数を使う
こちらの記事が参考になります。
root_dir: ${oc.env:PROJECT_ROOT}
output_dir: ${hydra:runtime.output_dir}
work_dir: ${hydra:runtime.cwd}
環境変数以外にも、実行時のログ保存ディレクトリなどの hydra 特有のパスにもアクセスできます。もちろん変更も可能なので、例えば指定した引数と同じ名前の出力ディレクトリを作ることなども可能です。
notebook で使う
with hydra.initialize(version_base=1.3, config_path="../configs"): CFG = hydra.compose( config_name="config.yaml", return_hydra_config=True, overrides=OVERRIDES, ) # use HydraConfig for notebook to use hydra job HydraConfig.instance().set_config(CFG)
OVERRIDES には transformers=default などのようにコマンドライン引数での設定に相当する部分を書けば OK です。
multirun
今回は特に触れてなかったですが、hydra の目玉機能の一つです。
複数の異なる設定で実行する場合に使います。例えば、学習率を変えた実験を行いたい場合などですね。直列での実行になりますが、Launcher を変更することで並列の実行を可能にします。
おわりに
hydra を使った設定管理について、よく使うシチュエーションと実際のコードを作ってみました。この記事に書いた使い方以外にもまだまだ色々なことができそうですし、やりたいことはほぼできると使ってみて感じました。
公式 github にもある Hydra Ecosystem なんかも参考になりそうです。
今回使用したコードはこちらにあります。
