Zenn
🙆‍♀️

Optunaを用いた分散環境でのハイパーパラメータ最適化：実践ガイド

2025/03/04に公開
機械学習
PyTorch
Optuna
中野哲平
tech

はじめに

Optunaは、機械学習のハイパーパラメータ最適化（HPO）に広く利用されているPythonライブラリです。特に、分散環境（マルチGPU・マルチノード）でのハイパーパラメータ探索 に対応しており、大規模なモデルのチューニングを効率的に行うことができます。

本記事では、

  • Optunaを分散環境で利用する方法
  • ハイパーパラメータ探索を効率化するテクニック
  • 実装時の落とし穴と注意点

を詳しく解説します。

1. 分散環境でのOptunaの利用方法

Optunaはマルチプロセス・マルチマシン環境で動作させることができ、特に以下の2つの方法で分散最適化を行うことが可能です。

方法1: Joblib による並列最適化（単一ノード内での並列化）

単一の計算ノード（1台のマシン）で複数のGPUを活用する場合、Joblib を使うことで複数のトライアルを並列に実行できます。

import optuna
from joblib import Parallel, delayed

def objective(trial):
    lr = trial.suggest_loguniform("lr", 1e-5, 1e-1)
    batch_size = trial.suggest_categorical("batch_size", [16, 32, 64])
    return train_model(lr, batch_size)

study = optuna.create_study(direction="maximize")
Parallel(n_jobs=4)(delayed(study.optimize)(objective, n_trials=10) for _ in range(4))
  • メリット：シンプルなコードで、複数のトライアルを並列実行できる。
  • デメリット：単一ノードに制限されるため、スケールアップが難しい。

方法2: RDB（データベース）を活用した分散最適化（マルチノード対応）

より大規模な分散最適化を行う場合、RDB（リレーショナルデータベース）をバックエンドに設定し、複数のマシンやGPUで同時に最適化を実行 できます。

  1. RDBサーバーを起動（例: SQLite, MySQL, PostgreSQL）
  2. 全てのノードで同じデータベースを指定してOptunaを実行
import optuna

# MySQLを使用する場合
storage = "mysql+pymysql://user:password@localhost/dbname"
study = optuna.create_study(study_name="distributed_hpo", storage=storage, direction="maximize")

# 各ノードで並列実行
study.optimize(objective, n_trials=100, n_jobs=4)
  • メリット：複数のマシンで同時にトライアルを実行でき、大規模な探索が可能。
  • デメリット：データベースの管理が必要で、通信コストが発生する。

どちらの方法を選ぶべきか？

  • ローカルで小規模に試す場合 → Joblib（単一ノード）
  • 本格的にスケールアップしたい場合 → RDBを利用（マルチノード）

2. 効率的なハイパーパラメータ探索のテクニック

1. Samplerを適切に選択する

Optunaには、複数のサンプリングアルゴリズムがあり、適切なものを選ぶことで探索効率を向上できます。

  • TPESampler()（デフォルト）：ベイズ最適化を用いた手法。一般的に良い選択肢。
  • RandomSampler()：完全ランダム探索。ベースラインとして有用。
  • GridSampler()：グリッドサーチ。少数のパラメータなら有効。
  • CmaEsSampler()：連続値の最適化に強い（進化戦略）。

例えば、TPE（Tree-structured Parzen Estimator）を使用する場合：

import optuna

study = optuna.create_study(sampler=optuna.samplers.TPESampler(), direction="maximize")
study.optimize(objective, n_trials=100)

2. 途中経過のトライアルを破棄する（Pruning）

学習の途中で明らかに良くないトライアルを早期終了することで、計算資源を節約できます。

from optuna.integration import PyTorchLightningPruningCallback

pruner = optuna.pruners.MedianPruner()
study = optuna.create_study(direction="maximize", pruner=pruner)

3. Optuna 分散学習時の注意点・落とし穴

1. データベースのボトルネック

分散最適化では、すべてのノードが中央のRDBにアクセスするため、データベースがボトルネックになる可能性 があります。

  • 解決策：PostgreSQLやMySQLを使用し、適切なインデックスを設定する。

2. GPUのリソース管理

  • 各トライアルが異なるGPUを使用するように管理しないと、複数のトライアルが同じGPUを使って競合し、クラッシュする可能性 があります。
  • 解決策：CUDA_VISIBLE_DEVICES を適切に設定し、各プロセスが異なるGPUを使用するようにする。
import os
os.environ["CUDA_VISIBLE_DEVICES"] = str(trial.number % 4)  # 4GPU構成の場合

3. study.optimize() の並列実行に関する問題

  • n_jobs を大きくしすぎると、CPU/GPUリソースの競合が発生。
  • 適切な n_trials を設定し、負荷分散を考慮する。

4. サーチスペースが広すぎると収束しない

  • suggest_uniform() で広すぎる範囲を指定すると、最適解に収束しにくくなる。
  • 適切な範囲を設定することが重要。
trial.suggest_uniform("dropout", 0.1, 0.5)  # 事前に適切な範囲を絞る

まとめ

Optunaを分散環境で活用することで、ハイパーパラメータ探索を劇的に高速化できます。

✅ 単一ノードでは Joblib を使用
✅ マルチノードでは RDB を活用
✅ 適切な Sampler を選択し、Pruning で計算リソースを節約
✅ GPU競合やデータベース負荷を避ける工夫が必要

適切な設定と戦略を用いることで、大規模なモデルの最適化がスムーズに進む ので、ぜひ実践してみてください！

Discussion