今回は、Hydra, MLflow, Optunaを使用して
機械学習モデルのハイパーパラメータ最適化を効率的に行う実装方法 を紹介します。
- Python Version 3.12.3
- Linux使用(deb系)
- 仮想環境にはvenvを使用
本記事では上記環境での実行を前提としています。
環境差異によるコマンドの違いについては扱わない点、ご了承ください。
機械学習モデルの開発では、次のような課題に直面することが多いです。
- 探索すべきハイパーパラメータ検討数が膨大
- 深層学習モデルハイパラだけでなく、データ前処理の条件も含めて検証したい
- 各条件での評価指標を自動で記録・可視化し、比較したい
本記事ではこれらの課題を解決するために、
- Hydra(設定管理)
- hydra-optuna-sweeper(ベイズ最適化)
- MLflow(実験管理・可視化)
を組み合わせた実装例を、コードとともに紹介します。
ディレクトリ構成
/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の先頭に次のように記述します。
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デコレータを渡し、
設定パラメータを受け取ります。
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から切り離しています。
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 など)を指します。
