PyTorchJobTuningSchedulerによる自動チューニング例
このドキュメントは、PyTorchJobTuningSchedulerを用いてKubernetesクラスタ上で自動的にハイパーパラメータチューニングを実行する例を説明します。
ここでは、intelligence/components/zenith-tune/examples/integration/kubernetesに含まれるMNISTの学習スクリプトに対して、learning rate、batch size、optimizerをアノテーション設定のみでチューニングする方法を示します。
前提条件
自動チューニング例を実行する前に、以下を準備してください：
- PyTorchJobを投入・取得可能なKubernetesクラスタ環境
- kubectlがインストールされていること
- ZenithTuneがインストールされたPython環境
アクセス可能なKubernetesクラスタを所有していない場合は、minikubeを用いてローカル環境にクラスタを構築することで自動チューニング例を実行できます。
# Minikubeクラスタの起動とコンテキスト設定
minikube start
kubectl config use-context minikube
# Kubeflow Training Operatorのインストール
kubectl apply --server-side -k "github.com/kubeflow/training-operator.git/manifests/overlays/standalone?ref=v1.8.1"
自動チューニングで使用するファイル
FAIBのintelligence/components/zenith-tune/examples/integration/kubernetesに含まれる、自動チューニングで使用するファイルについて説明します。
scheduler_example.py
このPythonスクリプトは、PyTorchJobTuningSchedulerを起動してアノテーションベースの自動チューニングを実行します。
スケジューラーは、zenith-tune/optimization-configアノテーションを持つジョブのみを対象とするよう設定されています：
from zenith_tune.integration.kubernetes import JobFilter, PyTorchJobTuningScheduler
# Only target jobs with zenith-tune/optimization-config annotation
job_filter = JobFilter(
annotations={"zenith-tune/optimization-config": None} # Key existence check
)
scheduler = PyTorchJobTuningScheduler(
submit_namespace="default",
job_filter=job_filter,
max_concurrent_tuning=1,
)
scheduler.run()
annotation_based_tuning.yaml
このYAMLファイルは、アノテーションベースのチューニング設定を含むPyTorchJobを定義しています。 学習スクリプトはConfigMapで定義され、コマンドライン引数からハイパーパラメータを受け取れるように設計されています。
学習スクリプトの主な特徴：
- learning rate、batch size、optimizerの設定をコマンドライン引数で受け取る
- 最終的なloss値を出力する（チューニングの目的関数）
PyTorchJobのアノテーション設定：
apiVersion: kubeflow.org/v1
kind: PyTorchJob
metadata:
generateName: annotation-based-tuning-
annotations:
zenith-tune/optimization-config: |
variables:
- name: "learning_rate"
type: "float"
range: [0.001, 0.1]
log: true
target_env: "LEARNING_RATE"
- name: "batch_size"
type: "int"
range: [16, 128]
step: 16
target_env: "BATCH_SIZE"
- name: "optimizer"
type: "categorical"
choices: ["sgd", "adam", "rmsprop"]
target_env: "OPTIMIZER"
objective:
name: "loss"
regex: "Final Loss: ([0-9]+\\.?[0-9]*)"
direction: "minimize"
n_trials: 5
spec:
pytorchReplicaSpecs:
Worker:
replicas: 1
restartPolicy: Never
template:
spec:
containers:
- name: pytorch
image: pytorch/pytorch:1.13.1-cuda11.6-cudnn8-runtime
command:
- sh
- -c
- |
python /scripts/train_mnist_tunable.py \
--learning-rate ${LEARNING_RATE:-0.01} \
--batch-size ${BATCH_SIZE:-32} \
--optimizer ${OPTIMIZER:-sgd} \
--epochs 2
自動チューニングの実行
1. スケジューラーの起動
# スケジューラーを実行
python scheduler_example.py
2. PyTorchJobの投入
スケジューラーが起動している状態で、チューニング対象のPyTorchJobをクラスタに投入します：
# ConfigMapとPyTorchJobをクラスタに投入
kubectl create -f annotation_based_tuning.yaml
3. チューニング状況の確認
# チューニングジョブの確認
kubectl get pytorchjobs
連続5回のチューニングジョブが実行され、ハイパーパラメータが特定されたら成功です。