機械学習の実験管理入門！パラメータをYAMLとシェルで自由に設定する方法

はじめに

機械学習で、パラメータ調整をする頻度はとても高いです。
今回、この記事を読むことで、

機械学習のパラメータを設定する環境を素早く作り、
変更をしたいときに、その都度簡単にシェルで指定できるようになります。

---

この記事は、以下のリポジトリとドキュメントを参考にして書いています。

動作環境

• Python 3.10.12
• WSL Ubuntu 22.04.5
• uv 0.9.8

1. プロジェクト作成

まずはPythonのプロジェクトを作成します

uv init tutorial
cd tutorial

初期状態は以下の状態になります

tutorial/ 
├── README.md
├── main.py
└── pyproject.toml

次に、設定管理のために omegaconf を追加します

uv add omegaconf

そして、src/config.yamlを手動で作成します。

tutorial/ 
├── README.md
├── main.py
├── pyproject.toml
├── src
│   └── config.yaml   ← 手動で作成
└── uv.lock

2. 設定用のYAMLファイルを作る

src/config.yamlの内容を、
機械学習のパラメータを例として記述していきます。

batch_size: 32
learning_rate: 0.001
epochs: 50
optimizer:
  type: adam
  weight_decay: 0.0001

YAML は階層的な構造にできるため、
以下のように、1つのファイルに
訓練時とテスト時を分けて書くこともできます

訓練時とテスト時のパラメータを、1つのファイルに記述した例

train:
  batch_size: 32
  learning_rate: 0.001
  epochs: 50
  optimizer:
    type: adam
    weight_decay: 0.0001
  shuffle: true

test:
  batch_size: 32
  shuffle: false
  model_path: ./checkpoints/best_model.pth

3. Python 側で読み込んで上書きする

main.pyの例が以下です。
argparseはPythonの標準ライブラリであるため追加する必要はありません。

以下の簡潔なコードで、YAMLファイルを読み込みつつ、
シェルで上書きが可能となります。

from omegaconf import OmegaConf
import argparse

parser = argparse.ArgumentParser()
parser.add_argument("--config", type=str, default="src/config.yaml")
parser.add_argument('opts', nargs=argparse.REMAINDER)
args = parser.parse_args()

cfg = OmegaConf.load(args.config)
cfg = OmegaConf.merge(cfg, OmegaConf.from_cli(args.opts))

print(OmegaConf.to_yaml(cfg))

4. 実行例

YAMLファイルの設定のまま実行

uv run python main.py

出力結果：

batch_size: 32
learning_rate: 0.001
epochs: 50
optimizer:
  type: adam
  weight_decay: 0.0001

CLI で上書き

uv run python main.py batch_size=64

出力結果：

batch_size: 64
learning_rate: 0.001
epochs: 50
optimizer:
  type: adam
  weight_decay: 0.0001

CLI で上書き（階層パラメータ）

uv run python main.py optimizer.type=mse

出力結果：

batch_size: 32
learning_rate: 0.001
epochs: 50
optimizer:
  type: mse
  weight_decay: 0.0001

おわりに

今回、
OmegaConfによって、YAML設定ファイルを使いつつ、
シェルで都度パラメータを変更できる環境の作り方について説明しました。

シェルで指定する際に、
この設定方法だと自由度が高いです。

しかし、タイプミスに怒ってくれることはないので、
十分に注意しましょう。

これで基本の環境はできたはずです！
いろいろな方法があると思うので、
ぜひコメントで教えてください！

記事を読んでくださりありがとうございました！