ハイパラ管理のすすめ -ハイパーパラメータを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に移行してみてはいかがでしょうか。

参考文献

• Hydra
• facebookresearch/hydra
• Hydra — A fresh look at configuration for machine learning projects