2
1

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

Hydra&MLflow&Optunaで機械学習モデルの最適化

2
Posted at

今回は、Hydra, MLflow, Optunaを使用して
機械学習モデルのハイパーパラメータ最適化を効率的に行う実装方法 を紹介します。


  • Python Version 3.12.3
  • Linux使用(deb系)
  • 仮想環境にはvenvを使用

本記事では上記環境での実行を前提としています。
環境差異によるコマンドの違いについては扱わない点、ご了承ください。


機械学習モデルの開発では、次のような課題に直面することが多いです。


  1. 探索すべきハイパーパラメータ検討数が膨大
  2. 深層学習モデルハイパラだけでなく、データ前処理の条件も含めて検証したい
  3. 各条件での評価指標を自動で記録・可視化し、比較したい

本記事ではこれらの課題を解決するために、

  • Hydra(設定管理)
  • hydra-optuna-sweeper(ベイズ最適化)
  • MLflow(実験管理・可視化)

を組み合わせた実装例を、コードとともに紹介します。

HydraMLflowOptuna (1).png


ディレクトリ構成

/lstm_analysis
  ├ conf/
  │  └ config.yaml
  ├ data/
  │  ├ dataset_load.py
  │  └ preprocess.py
  ├ models/
  │  ├ lstm.py
  │  └ xlstm.py
  ├ trainers/
  │  └ trainer.py
  └ main.py
  • main.py : エントリーポイント
  • conf/config.yaml : Hydra用設定ファイル
  • data/ : データ読み込み、前処理
  • models/ : モデル定義
  • trainers/ : 学習・評価処理

インストール

pip install hydra-core
pip install hydra-optuna-sweeper
pip install mlflow

Hydra optuna sweeperを用いたベイズ最適化

結論から言うと、Hydra Optuna Sweeperを用いることで、

python main.py --multirun

と実行することで、
指定したハイパーパラメータ空間に対してOptunaによるベイズ最適化を自動で実行できます。


HydraでOptunaを使用する方法

1. Sweeperの指定

config.yamlの先頭に次のように記述します。

config/config.yaml
defaults:
  - override hydra/sweeper: optuna

2. 通常の設定項目

dataset: something_v1
model: lstm

# --------------------
# データセット設定
# --------------------
datasets:
  something_v1:
    path: "something_datasets/something_v1.csv"
  something_v2:
    path: "something_datasets/something_v2.csv"

# --------------------
# モデル設定
# --------------------
models:
  lstm:
    type: lstm
    hidden_dim: 512
    num_layers: 2
    dropout: 0.2
  xlstm:
    type: xlstm
    hidden_dim: 512
    dropout: 0.2
# --------------------
# トレーニング設定
# --------------------
training:
  lr: 0.0002
  train_val_split: 0.2

3. Optuna Sweeperの設定

検証データのlossを最小化したいため、 direction: "minimize"を指定します。

hydra:
  sweeper:
    _target_: hydra_plugins.hydra_optuna_sweeper.optuna_sweeper.OptunaSweeper
    n_trials: 100          # optuna試行回数
    direction: "minimize"  # 評価指標最適化方向
    sampler:
      _target_: optuna.samplers.TPESampler
    params:                # ハイパラ空間
      training.lr: tag(log, interval(1e-6, 1e-3))
      models.lstm.hidden_dim: choice(128, 256, 512)
      models.lstm.num_layers: choice(1, 2, 3, 4)
      models.lstm.dropout: choice(0.1, 0.2, 0.3, 0.4)

ハイパーパラメータ空間の指定方法については、公式ドキュメントが参考になります。

Hydra使用方法

main関数に@hydra.mainデコレータを渡し、
設定パラメータを受け取ります。

main.py
import hydra
from omegaconf import DictConfig


@hydra.main(config_name="config", config_path="conf")
def main(cfg: DictConfig) -> None:
    
    # 1. dataset 名(文字列)
    dataset_name = cfg.dataset # something_v1

    # 2. datasets 辞書から該当データの path を取り出す
    dataset_path = cfg.datasets[dataset_name].path

    # 3. model設定取得
    model_cfg = cfg.models[cfg.model]
    model_type = model_cfg.type # lstm

    hidden_dim = model_cfg.hidden_dim # 512
    num_layers = model_cfg.num_layers # 2
    dropout = model_cfg.dropout # 0.2

    
    # 4. トレーニング設定取得
    lr = cfg.training.lr # 0.002
    train_val_split = cfg.training.train_val_split # 0.2

    # 学習処理 ...



if __name__ == "__main__":
    main()

MLflowによる学習ログの記録

MLflowを用いた学習ログの実装例です。
機械学習の実装にはPyTorchを使用し、コードの見通しを良くするため、モデル定義や学習過程はクラスやメソッドLSTMmodel, Trainer, trainer.fit()を定義してmain.pyから切り離しています。

main.py
import torch
import torch.nn as nn
import torch.optim as optim

# hydra & mlflow
import hydra
from omegaconf import DictConfig
import mlflow

# 独自クラス
from models.lstm import LSTMModel
from trainers.trainer import Trainer

    # 学習データ読み込みなど

    mlflow.set_experiment(f"{model_type}_exp_{dataset_name}") # lstm_exp_something_v1

    # Optunaで最適化する指標保存用
    metric_return = float('inf')
    
    with mlflow.start_run():
        mlflow.log_param("hidden_dim", hidden_dim)
        mlflow.log_param("num_layers", num_layers)
        mlflow.log_param("dropout", dropout)
        mlflow.log_param("lr", lr)
        mlflow.log_param("train_val_split", train_val_split)
        
        model = LSTMModel(
            input_dim=12,
            hidden_dim=hidden_dim,
            num_layers=num_layers,
            output_dim=4,
            dropout=dropout
        )

        # ====== 損失関数と最適化関数 ======
        criterion = nn.MSELoss()
        optimizer = optim.Adam(model.parameters(), lr=lr)

        
        # ====== 学習・検証 =====
        trainer = Trainer(
            model=model,
            optimizer=optimizer,
            criterion=criterion,
            )

        # ====== メイン学習ルーチン ======
        train_loss, valid_loss = trainer.fit( ... ) # エポック数分学習し評価指標を返す

        mlflow.log_metric("train_loss", train_loss)
        mlflow.log_metric("valid_loss", valid_loss)

        metric_return = valid_loss
    # end with

    return metric_return # optunaが最適化する指標をmain関数の返り値とする


if __name__ == "__main__":
    main()
  • mlflow.set_experiment() : 比較する実験ログをセット
  • mlflow.start_run() : ハイパラや評価指標の記録をスタート
  • mlflow.log_param() : ハイパラを記録
  • mlflow.log_metric() : 評価指標を記録

main関数の返り値がOptunaの最適化対象になります。


コマンドラインによる実行

次のように記述することで、指定したdatasetについてoptunaによるベイズ最適化を走らせることができます。

python main.py --multirun dataset=something_v1

MLflow UI

mlflow ui --port 5000

ブラウザでhttp://127.0.0.1:5000にアクセスすると、解析結果を一覧で確認することができます。

参考資料

補足

pythonのデコレータとは

デコレータを使用した関数の挙動はそのままに、前後に特定の操作を行わせる機能。(関数をデコレーションする。)
@hydra.mainは、設定管理やマルチランを可能にする修飾を行っています

Metric

モデルの性能を数値で評価する指標(loss, accuracy など)を指します。

2
1
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
2
1

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?