ハイパラ管理のすすめ -ハイパーパラメータをHydra+MLflowで管理しよう-
機械学習をやっている人なら誰もが遭遇したであろうこの光景
(※写真はPyTorchのLanguage ModelのExampleより)
Pythonのargparseでシェルから引数を受け取りPythonスクリプト内でパラメータに設定するパターンは、記述が長くなりがちな上、どのパラメータがmodel/preprocess/optimizerのものなのか区別がつきにくく見通しが悪いといった課題があります。
私は実験用のパラメータ類は全てYAMLに記述して管理しています。
YAMLで記述することでパラメータを階層立てて構造的に記述することができ、パラメータの見通しがぐっとよくなります。
preprocess: min_df: 3 max_df: 1 replace_pattern: \d+ model: hidden_size: 256 dropout: 0.1 optimizer: algorithm: Adam learning_rate: 0.01 norm: 0.001
パラメータチューニングの際には、シェルスクリプトからyqコマンドで書き換えながらPythonスクリプトに流すという運用をしていたのですが、yqコマンドでがちゃがちゃ書き直しているうちにデフォルト値が分からなくなるという悩みがありました。
YAMLによるパラメータ管理のベストプラクティスを模索している折に、Hydraというツールが登場したので、家の実験管理周りをHydraを使って整理してみました。
Hydraとは
HydraはFacebook Researchが提供している設定ファイルを管理しやすくするためのツールです。
様々な設定をYAML形式で記述し、そのYAMLの設定群を簡単にPythonスクリプト内に流し込むことに主眼を置いているツールであり、ExampleにはDatabaseの設定があるなど機械学習以外の用途での使用も想定しているツールです。
HydraによるYAMLの読み込み
以下のようにYAMLファイルにパラメータを設定し、PythonスクリプトでHydraのデコレータを付与した関数を用意することでDictの形式でパラメータを読み込むことができるようになります。
config.yaml
db: driver: postgresql pass: drowssap timeout: 20 user: postgre_user
my_app.py
@hydra.main(config_path='config.yaml') def my_app(cfg): print(cfg.pretty())
$ python my_app.py
db:
driver: postgresql
pass: drowssap
timeout: 20
user: postgre_user
コマンドラインでYAMLのパラメータをkey=valueの形で渡すと、対象の値を書き換えてPythonスクリプトに持ち込むことができます。もちろん元のYAMLファイルには影響はありません。
$ python my_app.py db.user=ymym db.pass=3412 db: driver: postgresql pass: 3412 timeout: 20 user: ymym
複数のYAMLファイルの管理
Hydraでは設定ファイルを複数のYAMLファイルに分割して運用することも想定しています。
例えば、NNとLightGBMのパラメータを別々のYAMLファイルに記述して使用するモデルをハイパラに設定してそれに応じて対応するモデルのYAMLを読み込といった感じです。
nn.yaml
model: layers: 3 dropout: 0.5
lightgbm.yaml
model: max_depth: 10 learning_rate: 0.01
以下のようにディレクトリを切ってYAMLを配置して、どの設定ファイルを読み込むかを config.yaml で制御します。
├── conf │ ├── config.yaml │ └── model │ ├── lightgbm.yaml │ └── nn.yaml └── my_app.py
config.yaml
defaults: - model: nn
$ python my_app.py model: layers: 3 dropout: 0.5
Hydraの出力ディレクトリ
HydraはPythonスクリプトが最終的にどんなYAMLファイルの内容で実行されたかを出力ディレクトリ(デフォルトではoutputs/)を生成して保管してくれます。
├── .hydra │ ├── config.yaml │ ├── hydra.yaml │ └── overrides.yaml └── my_app.log
この出力ディレクトリには少し注意が必要で、Pythonスクリプトでhydraのデコレータをつけた関数の中ではcwdがこの出力ディレクトリになってしまいます。
Pythonコードの中で pd.read_csv('data/train.csv')といったファイル読み込みを使用とするとcwdの違いから事故ることが多いので、hydraが用意してくれている関数を使ってオリジナルのプロジェクトルートのパスを取得するとよいでしょう。
import os from omegaconf import DictConfig import hydra @hydra.main() def my_app(cfg: DictConfig) -> None: print(f'Current working directory: {os.getcwd()}') print(f'Orig working directory : {hydra.utils.get_original_cwd()}') print(f'to_absolute_path("foo") : {hydra.utils.to_absolute_path("foo")}') print(f'to_absolute_path("/foo") : {hydra.utils.to_absolute_path("/foo")}') >>>Current working directory: /home/user/workspace/hydra-exp/outputs/2020-02-09/02-29-26 >>>Orig working directory : /home/user/workspace/hydra-exp >>>to_absolute_path("foo") : /home/user/workspace/hydra-exp/foo >>>to_absolute_path("/foo") : /foo
Hydra + MLflowでパラメータ/実験を管理する
では、機械学習の実験に対して「YAMLで記述したハイパーパラメータの読み込みとグリッドサーチにHydraを」「どのパラメータで実験しどんな結果になったかの記録をMLflow」で行います。
今回も題材は例によってLivedoorのニュースコーパスのテキスト分類です。
まずはデータの読み込み、加工等の諸々の関数を定義します。
# AllenNLP用に文章からInstanceを生成する def text_to_instance(word_list, label): tokens = [Token(word) for word in word_list] word_sentence_field = TextField(tokens, {"tokens": SingleIdTokenIndexer()}) fields = {"tokens": word_sentence_field} if label is not None: label_field = LabelField(label, skip_indexing=True) fields["label"] = label_field return Instance(fields) def load_dataset(path, dataset): if dataset not in ['train', 'val', 'test']: raise ValueError('"dataset" parametes must be train/val/test') data, labels = pd.read_csv(f'{path}/{dataset}.csv'), pd.read_csv(f'{path}/{dataset}_label.csv', header=None, squeeze=True) return data, labels def preprocess(X, y, preprocessor=None): if preprocessor is None: preprocessor = Preprocessor() preprocessor\ .stack(ct.text.UnicodeNormalizer())\ .stack(ct.Tokenizer("ja"))\ .fit(X['article']) processed = preprocessor.transform(X['article']) dataset = [text_to_instance([token.surface for token in document], int(label)) for document, label in zip(processed, y)] return dataset, preprocessor
次にハイパーパラメータを記述するYAMLファイルです。
config.yaml
# word embeddingに関するハイパーパラメータ w2v: model_name: all vocab_size: 32000 norm: 2 # モデルに関するパラメータ model: hidden_size: 256 dropout: 0.5 # 実験時に使用するパラメータ training: batch_size: 32 learning_rate: 0.01 epoch: 30 patience: 3
今回はYAMLは分割せずひとつのファイルにすべて記述しています。
個人的にはYAMLを細く分割しすぎると変更忘れや修正がおっくうになるので、それほど複雑でなければ単一のYAMLにまとめて記述してしまった方が良いと思います。
続いてTrain&Testの関数です。
# 学習 def train(train_dataset, val_dataset, cfg): # Vocabularyを生成 VOCAB_SIZE = cfg.w2v.vocab_size vocab = Vocabulary.from_instances(train_dataset + val_dataset, max_vocab_size=VOCAB_SIZE) BATCH_SIZE = cfg.training.batch_size # パディング済みミニバッチを生成してくれるIterator iterator = BucketIterator(batch_size=BATCH_SIZE, sorting_keys=[("tokens", "num_tokens")]) iterator.index_with(vocab) # 東北大が提供している学習済み日本語 Wikipedia エンティティベクトルを使用する # http://www.cl.ecei.tohoku.ac.jp/~m-suzuki/jawiki_vector/ model_name = cfg.w2v.model_name norm = cfg.w2v.norm cwd = hydra.utils.get_original_cwd() params = Params({ 'embedding_dim': 200, 'padding_index': 0, 'pretrained_file': os.path.join(cwd, f'embs/jawiki.{model_name}_vectors.200d.txt'), 'norm_type': norm}) token_embedding = Embedding.from_params(vocab=vocab, params=params) HIDDEN_SIZE = cfg.model.hidden_size dropout = cfg.model.dropout word_embeddings: TextFieldEmbedder = BasicTextFieldEmbedder({"tokens": token_embedding}) encoder: Seq2SeqEncoder = PytorchSeq2SeqWrapper(nn.LSTM(word_embeddings.get_output_dim(), HIDDEN_SIZE, bidirectional=True, batch_first=True)) model = ClassifierWithAttn(word_embeddings, encoder, vocab, dropout) model.train() USE_GPU = True if USE_GPU and torch.cuda.is_available(): model = model.cuda(0) LR = cfg.training.learning_rate EPOCHS = cfg.training.epoch patience = cfg.training.patience if cfg.training.patience > 0 else None optimizer = optim.Adam(model.parameters(), lr=LR) trainer = Trainer( model=model, optimizer=optimizer, iterator=iterator, train_dataset=train_dataset, validation_dataset=val_dataset, patience=patience, cuda_device=0 if USE_GPU else -1, num_epochs=EPOCHS ) metrics = trainer.train() logger.info(metrics) return model, metrics def test(test_dataset, model, writer): # 推論 model.eval() with torch.no_grad(): predicted = [model.forward_on_instance(d)['logits'].argmax() for d in tqdm(test_dataset)] # Accuracyの計算 target = np.array([ins.fields['label'].label for ins in test_dataset]) predict = np.array(predicted) accuracy = accuracy_score(target, predict) # Precision/Recallの計算 macro_precision = precision_score(target, predict, average='macro') micro_precision = precision_score(target, predict, average='micro') macro_recall = recall_score(target, predict, average='macro') micro_recall = recall_score(target, predict, average='micro') # MLflowに記録 writer.log_metric('accuracy', accuracy) writer.log_metric('macro-precision', macro_precision) writer.log_metric('micro-precision', micro_precision) writer.log_metric('macro-recall', macro_recall) writer.log_metric('micro-recall', micro_recall) model.cpu() writer.log_torch_model(model)
ここで出てくるwriterというインスタンスはMLflowのClientをラップしてログの記録やArtifactの保存を行うクラスのインスタンスです。
with mlflow.start_run():のブロック外でもMLflowを使う場面があり、Run IDを引き回さないといけないためラッパークラスを作っています。
class MlflowWriter(): def __init__(self, experiment_name, **kwargs): self.client = MlflowClient(**kwargs) try: self.experiment_id = self.client.create_experiment(experiment_name) except: self.experiment_id = self.client.get_experiment_by_name(experiment_name).experiment_id self.run_id = self.client.create_run(self.experiment_id).info.run_id def log_params_from_omegaconf_dict(self, params): for param_name, element in params.items(): self._explore_recursive(param_name, element) def _explore_recursive(self, parent_name, element): if isinstance(element, DictConfig): for k, v in element.items(): if isinstance(v, DictConfig) or isinstance(v, ListConfig): self._explore_recursive(f'{parent_name}.{k}', v) else: self.client.log_param(self.run_id, f'{parent_name}.{k}', v) elif isinstance(element, ListConfig): for i, v in enumerate(element): self.client.log_param(self.run_id, f'{parent_name}.{i}', v) def log_torch_model(self, model): with mlflow.start_run(self.run_id): pytorch.log_model(model, 'models') def log_param(self, key, value): self.client.log_param(self.run_id, key, value) def log_metric(self, key, value): self.client.log_metric(self.run_id, key, value) def log_artifact(self, local_path): self.client.log_artifact(self.run_id, local_path) def set_terminated(self): self.client.set_terminated(self.run_id)
最後にHydraのデコレータを付与したmain関数です。
データをローカルのcsvから読み込むため、Hydraのutilを使ってプロジェクトルートのパスを取得しています。
@hydra.main(config_path='config.yaml') def main(cfg: DictConfig): # https://medium.com/pytorch/hydra-a-fresh-look-at-configuration-for-machine-learning-projects-50583186b710 cwd = hydra.utils.get_original_cwd() train_X, train_y = load_dataset(os.path.join(cwd, 'data'), 'train') val_X, val_y = load_dataset(os.path.join(cwd, 'data'), 'val') test_X, test_y = load_dataset(os.path.join(cwd, 'data'), 'test') train_dataset, preprocessor = preprocess(train_X, train_y) val_dataset, preprocessor = preprocess(val_X, val_y, preprocessor) test_dataset, preprocessor = preprocess(test_X, test_y, preprocessor) EXPERIMENT_NAME = 'livedoor-news-hydra-exp' writer = MlflowWriter(EXPERIMENT_NAME) writer.log_params_from_omegaconf_dict(cfg) model, metrics = train(train_dataset, val_dataset, cfg) test(test_dataset, model, writer) # Hydraの成果物をArtifactに保存 writer.log_artifact(os.path.join(os.getcwd(), '.hydra/config.yaml')) writer.log_artifact(os.path.join(os.getcwd(), '.hydra/hydra.yaml')) writer.log_artifact(os.path.join(os.getcwd(), '.hydra/overrides.yaml')) writer.log_artifact(os.path.join(os.getcwd(), 'main.log')) writer.set_terminated() if __name__ == '__main__': main()
HydraにはMulti-runという機能があり、これはコマンドラインから呼ぶ際にパラメータのkey=valueでvalue値をカンマ区切りで記述し-mオプションをつけると、全パラメータの組み合わせを実行してくれるというものです。
また出力は各パラメータの組み合わせのたびに保存されるので、この機能を使ってパラメータのグリッドサーチを行うことができます。
$ python main.py w2v.model_name=all,entity,word model.hidden_size=32,64,128,256 training.learning_rate=0.01,0.005 -m
上記を実行すれば各実験の内容がMLflow上に記録されます。
Hydraでグリッドサーチした結果をMLflowに記録しておけば、実験結果の比較も容易です。
以下は各実験のパラメータを表示しながら、Accuracyをプロットしているところです。
まとめ
今回の記事ではFacebook Researchが開発している設定管理ツールのHydraの使い方と、Hydra+MLflowでハイパーパラメータの入出力を管理するやり方を紹介しました。
argparseを使ってパラメータ入力を行うのと比べて、YAMLでのパラメータ管理は見通しがよくHydraと組み合わせることで設定をいじりながらPythonと組み合わせることも簡単になります。
これを機にハイパーパラメータの管理をYAML+Hydraに移行してみてはいかがでしょうか。
