Azure Machine Learning CLI v2 でハイパーパラメーターチューニングジョブを投げてみる

はじめに

Azure Machine Learning というのは ML を支える Azure の便利サービスです。AWS であれば SageMaker 、 GCP であれば Vertex AI あたりと似た立ち位置になります。

CUDA セットアップ済みで即 GPU 使える VM やクラスターを GUI ポチポチで建てたり、パラメーター設定やコードのスナップショット、モデルの保存 (つまり実験管理) ができたりします。

さらには僕の推し実験管理ツールであるところの MLflow と互換性があり、定期的に Qiita などで「MLflow は良いゾ、Azure ML を MLflow-as-a-Service として使うと良いゾ」という記事を書いてたりします。無駄に電気を GPU に通して熱に変える前に実験管理しましょうね。

そんな Azure Machine Learning ですが、5月末に極めて大規模なアップデートが行われました。どれくらい大規模だったかは以下ページの「Data & AI: Machine Learning」のところの長さを見てもらえれば何となくわかるかと思います。僕の主戦場は Azure Machine Learning ということで現在必死にキャッチアップ中です。

そんな Azure ML 大規模アップデートの1つとして、Azure ML CLI v2 が GA となりました。これは az ml コマンドとジョブ定義を記述した yaml を使用して、 bash などから Azure Machine Learning に対してジョブ実行命令を出すことができる CLI です。

今回はこれを使ってハイパーパラメーターチューニングジョブを実行してみます。

事前準備

前提条件

Azure Machine Learning ワークスペースと conda が必要です。以下にドキュメントリンクを貼っておきます。

CLI v2 の使用準備

まずお手元の PC に Azure CLI と az ml コマンドを使うための extension をインストールします。

Azure CLI のインストール手順は環境ごとに異なります。

Windows にインストールするとき、僕は winget を使っています。

winget install Microsoft.AzureCLI

続いて ml extension をインストールします。

az extension add -n ml -y

ジョブの実行

スクリプトの準備

模擬用のデータセットと学習用のコンピューティングクラスターを構築するスクリプトを用意して実行しておきます。

prepare.py

import copy
from datetime import datetime

import pandas as pd
from azureml.core import Dataset, Datastore, Workspace
from azureml.core.compute import AmlCompute, ComputeTarget
from azureml.core.compute_target import ComputeTargetException
from azureml.opendatasets import NycTlcGreen
from dateutil.relativedelta import relativedelta

def register_dataset(ws: Workspace) -> None:
    dataset_name = "nyc_taxi_dataset"
    try:
        dataset = Dataset.get_by_name(ws, dataset_name)
        df = dataset.to_pandas_dataframe()

    except Exception:
        raw_df = pd.DataFrame([])
        start = datetime.strptime("1/1/2015", "%m/%d/%Y")
        end = datetime.strptime("1/31/2015", "%m/%d/%Y")

        for sample_month in range(3):
            temp_df_green = NycTlcGreen(
                start + relativedelta(months=sample_month), end + relativedelta(months=sample_month)
            ).to_pandas_dataframe()
            raw_df = raw_df.append(temp_df_green.sample(2000))

        raw_df.head(10)

        df = copy.deepcopy(raw_df)

        columns_to_remove = [
            "lpepDropoffDatetime",
            "puLocationId",
            "doLocationId",
            "extra",
            "mtaTax",
            "improvementSurcharge",
            "tollsAmount",
            "ehailFee",
            "tripType",
            "rateCodeID",
            "storeAndFwdFlag",
            "paymentType",
            "fareAmount",
            "tipAmount",
        ]
        for col in columns_to_remove:
            df.pop(col)

        df = df.query("pickupLatitude>=40.53 and pickupLatitude<=40.88")
        df = df.query("pickupLongitude>=-74.09 and pickupLongitude<=-73.72")
        df = df.query("tripDistance>=0.25 and tripDistance<31")
        df = df.query("passengerCount>0 and totalAmount>0")

        df["lpepPickupDatetime"] = df["lpepPickupDatetime"].map(lambda x: x.timestamp())

        datastore = Datastore.get_default(ws)
        dataset = Dataset.Tabular.register_pandas_dataframe(df, datastore, dataset_name)

    df.head(5)


def create_compute_cluster(ws: Workspace) -> None:
    cpu_cluster_name = "cpu-cluster"

    try:
        cpu_cluster = ComputeTarget(workspace=ws, name=cpu_cluster_name)
        print("Found existing cluster, use it.")
    except ComputeTargetException:
        compute_config = AmlCompute.provisioning_configuration(vm_size="STANDARD_D2_V2", max_nodes=8)
        cpu_cluster = ComputeTarget.create(ws, cpu_cluster_name, compute_config)

    cpu_cluster.wait_for_completion(show_output=False)


def main() -> None:
    subscription_id = "subscription_id"
    resource_group = "resource_group_name"
    workspace_name = "ml_workspace_name"

    ws = Workspace(
        workspace_name=workspace_name,
        subscription_id=subscription_id,
        resource_group=resource_group,
    )

    register_dataset(ws)
    create_compute_cluster(ws)


if __name__ == "__main__":
    main()

register_dataset関数は Azure のオープンデータセットから NycTlcGreen をダウンロードし、加工を施して Azure Machine Learning に Dataset として登録する関数です。create_compute_clusterは cpu-cluster という名前のクラスターを構築する関数です。パラメーターチューニングのためできるだけ並列数を稼ぎたいので、8ノード構成としています。 (特に動いていないときは勝手に停止するので課金事故も安心です)

続いて機械学習を実際に行うスクリプトを用意していきます。

train.py

import argparse
from typing import TypedDict

import lightgbm as lgb
import mlflow
import pandas as pd
from azureml.core import Dataset, Run
from sklearn.model_selection import train_test_split


class GetArgsOutput(TypedDict):
    input_dataset_name: str
    boosting_type: str
    metric: str
    learning_rate: float
    num_leaves: int
    min_data_in_leaf: int
    num_iteration: int


class LoadDatasetOutput(TypedDict):
    x_train: pd.DataFrame
    y_train: pd.DataFrame
    x_test: pd.DataFrame
    y_test: pd.DataFrame


def get_args() -> GetArgsOutput:
    # 引数取得
    parser = argparse.ArgumentParser()
    parser.add_argument("--input_dataset_name", type=str)
    parser.add_argument("--boosting_type", type=str, default="gbdt")
    parser.add_argument("--metric", type=str, default="rmse")
    parser.add_argument("--learning_rate", type=float, default=0.1)
    # FIXME: cli v2 の都合で int の探索空間を設定できないため float を渡してキャスト
    parser.add_argument("--num_leaves", type=float, default=10)
    # FIXME: cli v2 の都合で int の探索空間を設定できないため float を渡してキャスト
    parser.add_argument("--min_data_in_leaf", type=float, default=1)
    parser.add_argument("--num_iteration", type=int, default=100)

    args = parser.parse_args()

    params: GetArgsOutput = {
        "input_dataset_name": args.input_dataset_name,
        "boosting_type": args.boosting_type,
        "metric": args.metric,
        "learning_rate": args.learning_rate,
        "num_leaves": int(args.num_leaves),
        "min_data_in_leaf": int(args.min_data_in_leaf),
        "num_iteration": args.num_iteration,
    }

    return params


def load_dataset(input_dataset_name: str) -> LoadDatasetOutput:
    run = Run.get_context()
    ws = run.experiment.workspace

    dataset = Dataset.get_by_name(ws, input_dataset_name)
    df = dataset.to_pandas_dataframe()

    train, test = train_test_split(df, test_size=0.2, random_state=1234)

    x_train = train[train.columns[train.columns != "totalAmount"]]
    y_train = train["totalAmount"]

    x_test = test[test.columns[test.columns != "totalAmount"]]
    y_test = test["totalAmount"]

    output: LoadDatasetOutput = {"x_train": x_train, "y_train": y_train, "x_test": x_test, "y_test": y_test}

    return output


def train_lgb_model(args: GetArgsOutput, datasets: LoadDatasetOutput) -> lgb.Booster:
    train_dataset = lgb.Dataset(datasets["x_train"], datasets["y_train"])
    eval_dataset = lgb.Dataset(datasets["x_test"], datasets["y_test"], reference=train_dataset)

    mlflow.lightgbm.autolog(registered_model_name="nyc_taxi_regressor_lightgbm")

    params = {
        "boosting_type": args["boosting_type"],
        "metric": args["metric"],
        "learning_rate": args["learning_rate"],
        "num_leaves": args["num_leaves"],
        "min_data_in_leaf": args["min_data_in_leaf"],
        "num_iteration": args["num_iteration"],
        "task": "train",
        "objective": "regression",
    }
    mlflow.log_params(params)

    gbm = lgb.train(params, train_dataset, num_boost_round=50, valid_sets=eval_dataset, early_stopping_rounds=10)

    target_metric = "best_" + str(params["metric"])
    mlflow.log_metric(target_metric, gbm.best_score["valid_0"]["rmse"])

    return gbm


def main() -> None:
    args = get_args()
    datasets = load_dataset(args["input_dataset_name"])
    train_lgb_model(args=args, datasets=datasets)


if __name__ == "__main__":
    main()

register_dataset関数で登録したデータセットを pandas のデータフレームとして取得し、LightGBM で totalAmount を予測する回帰タスクです。このとき mlflow の autolog を使うことで各種メトリクスを自動で記録しています。

ポイントとなるのは調整対象となるパラメーター (num_leaves など) をスクリプトの引数として受け取っている点と、モデルのベストスコア (best_rmse) をmlflow.log_metricで記録している点です。パラメーターを外から渡して、そのパラメーターのときのスコアを記録することができる構成になっています。これにより、パラメーターの探索空間を定義し、特定のメトリックを最大化もしくは最小化するハイパーパラメーターチューニングが実行できるようになります。

型ヒントもつけています。VSCode が色々サジェストしてくれたり変な値の受け渡しを弾けるようになるのでお勧めです。入力が pandas なのか list なのか ndarray なのかはたまた torch.tensor なのか分からなくなるというような事態も防げるようになるので多人数開発のときに同僚に優しいコードになりますし、1人で開発しているとしても1か月前の自分とか実質他人みたいなものなので嬉しいことは多いと思います。

ジョブ定義

sweep job でハイパーパラメーターチューニングを実行することができます。定義を書いていきます。

search_hyperparameter.yml

$schema: https://azuremlschemas.azureedge.net/latest/sweepJob.schema.json
type: sweep
trial:
  code: src
  command: >-
    python train.py 
    --input_dataset_name ${{inputs.dataset_name}}
    --boosting_type ${{inputs.boosting_type}}
    --metric ${{inputs.metric}}
    --num_iteration ${{inputs.num_iteration}}
    --learning_rate ${{search_space.lr}}
    --num_leaves ${{search_space.num_leaves}}
    --min_data_in_leaf ${{search_space.min_data_in_leaf}}
  environment:
    image: mcr.microsoft.com/azureml/openmpi4.1.0-ubuntu20.04:20220504.v1
    conda_file: environment.yml

inputs:
  dataset_name: nyc_taxi_dataset
  boosting_type: gbdt
  metric: rmse
  num_iteration: 300

compute: azureml:cpu-cluster

sampling_algorithm: bayesian
search_space:
  lr:
    type: uniform
    min_value: 0.001
    max_value: 0.5
  num_leaves:
    type: uniform
    min_value: 10
    max_value: 200
  min_data_in_leaf:
    type: uniform
    min_value: 1
    max_value: 100

objective:
  goal: minimize
  primary_metric: best_rmse

limits:
  max_total_trials: 100
  max_concurrent_trials: 8
  timeout: 7200

experiment_name: job_nyc_taxi_regression_cli_v2_parameter_tuning
description: Hyper-parameter tuning job

上から順に見ていきます。

trial

code にはスクリプトを収めたディレクトリを指定する様子です。今回のジョブの場合は root ディレクトリにsearch_hyperparameter.yml を配した上で、 src/train.py を配置しています。

command ではコマンドを定義しています。今回は train.py という Python スクリプトを引数と共に実行する形です。${{hoge}}は後に登場する変数です。

environment は inline environment という形式をとっています。search_hyperparameter.yml と同じ階層に配置した以下 environment.yml を使用して、その場限りの Azure ML Environment を定義して使用する形です。ジョブを実行するための依存パッケージなどを定義しています。

environment.yml

name: env-ml-reposiotry-azureml-cli-v2
channels:
  - conda-forge
  - defaults
dependencies:
  - python=3.8
  - lightgbm=3.3.2
  - pandas=1.4.2
  - numpy=1.22.4
  - scikit-learn=1.1.1
  - mlflow=1.26.1
  - pip:
      - azureml-core==1.42.0
      - azureml-mlflow
      - azureml-opendatasets

inputs

ここでは変数を定義しています。例えばdataset_nameは trial ブロック内の${{inputs.dataset_name}}に代入されます。変数をまとめてここに定義することで、見通しが良くなります。

compute

prepare.py で作成したコンピューティングクラスターの名前を指定しています。

search_space

sampling_algorithm より先にこちらについて説明します。

search_space ではパラメーターの探索範囲を指定します。値と分布 (どういう偏りを持って値を選び出すか)を指定しています。

2022/06/07 現在、手元で試した限り整数を扱える randint や loguniform は動作せず、uniform と choice のみがうまく動きました。まだ機能強化の最中なのかもしれません。このため、スクリプトでは本来 int とすべきところを float で受け取って int 型にキャストするようにしています。

sampling_algorithm

これはパラメーターチューニングを行う際に、search_space で定義された範囲からどのようにパラメーターをサンプリングするかを指示する項目です。グリッドサーチか完全なランダムか、もう少し賢くベイズでやるかを指定できる様子です。

objective

primary_metric は実際に最適化するメトリックです。mlflow.log_metric("best_rmse", 0.2)のように、ここで指定した名前のメトリックをスクリプト内から記録しておく必要があります。

goal は具体的に primary_metric を最小化したいのか最大化したいのかを指示します。

limits

max_total_trials で試行回数の上限を決め、max_concurrent_trials で並列数を決め、timeout でジョブの実行上限時間を決めます。今回 max_concurrent_trials はコンピューティングクラスターのノード数に合わせて設定しています。

実行

ジョブを実行します。root 直下に search_hyperparameter.yml と environment.yml、src ディレクトリの中に train.py が収められているか確認してください。

az ml job create -f ./search_hyperparameter.yml -g <resouce_group_name> -w <ml_workspace_name>

結果

最良スコアは以下パラメーターに対して rmse 3.629 でした。

{"lr": 0.21120136633159164, "num_leaves": 71.59948662319685, "min_data_in_leaf": 46.53258476152086}

終わりに

Azure ML CLI v2 を使用してハイパーパラメーターチューニングを行うジョブを静的に定義してコマンドベースで実行することができました。

yaml 書いてコマンド叩いてジョブ実行という流れが良いですね。Python SDK でがりがり色々書くよりこちらの方が個人的には好きです。ipynb みたいなところでジョブ定義されてるより、.pyと.yaml の方が github でレビューもしやすいですし、見通しも良いですし。

深層学習モデルのようにモデルの構造を色々いじるときに、この CLI によるジョブ実行の仕組みを Github Acions に組み込めば楽に実験できそうだなぁとぼんやり思いました。実行時に src ディレクトリ配下を丸ごとアップロードしてる様子なので、そのときのコードやパラメーター設定含めて全部漏れなく記録できる点も非常に良いです。別記事でやってみようと思います。