在 AzureML 上使用 RAPIDS 进行训练和超参数调优

选择一组最优的超参数是一项艰巨的任务，特别是对于像 XGBoost 这样有许多超参数需要调优的算法。

在本 notebook 中，我们将展示如何在Azure 机器学习 (AzureML) 服务上并行运行多个训练作业，从而加快超参数优化速度。

前提条件

查看文档

创建 Azure ML 工作区，然后按照Microsoft Azure 机器学习中的说明启动带有 RAPIDS 的 Azure ML 计算实例。

一旦您的实例运行并可以访问 Jupyter，请保存此 notebook 并运行其中的单元格。

访问文档 >>

初始化工作区

初始化 MLClient 类以处理您在前提条件步骤中创建的工作区。

您可以手动提供工作区详细信息或调用 MLClient.from_config(credential, path) 从存储在 config.json 中的详细信息创建工作区对象。

from azure.ai.ml import MLClient
from azure.identity import DefaultAzureCredential

# Get a handle to the workspace.
#
# Azure ML places the workspace config at the default working
# directory for notebooks by default.
#
# If it isn't found, open a shell and look in the
# directory indicated by 'echo ${JUPYTER_SERVER_ROOT}'.
ml_client = MLClient.from_config(
    credential=DefaultAzureCredential(),
    path="./config.json",
)

从数据存储 URI 访问数据

在此示例中，我们将使用 2000 万行的航空公司数据集。下面的数据存储 URI 引用了一个包含 parquet 文件的存储位置（路径）。

datastore_name = "workspaceartifactstore"
dataset = "airline_20000000.parquet"

# Datastore uri format:
data_uri = f"azureml://subscriptions/{ml_client.subscription_id}/resourcegroups/{ml_client.resource_group_name}/workspaces/{ml_client.workspace_name}/datastores/{datastore_name}/paths/{dataset}"

print("data uri:", "\n", data_uri)

创建 AML 计算

您需要创建一个 Azure ML 托管计算目标（AmlCompute）作为训练模型的环境。

此 notebook 将使用 10 个节点进行超参数优化，您可以根据所需区域的可用配额修改 max_instances。与其他 Azure ML 服务类似，AmlCompute 存在限制，这篇文章详细介绍了默认限制以及如何请求更多配额。

size 描述了将在集群中使用的虚拟机类型和大小。请参阅 RAPIDS 文档中的“系统要求”（链接）和 Azure 文档中的“GPU 优化虚拟机大小”（链接）以确定实例类型。

让我们创建一个由 Standard_NC12s_v3（Tesla V100）GPU VM 组成的 AmlCompute 集群

from azure.ai.ml.entities import AmlCompute
from azure.ai.ml.exceptions import MlException

# specify aml compute name.
target_name = "rapids-cluster"

try:
    # let's see if the compute target already exists
    gpu_target = ml_client.compute.get(target_name)
    print(f"found compute target. Will use {gpu_target.name}")
except MlException:
    print("Creating a new gpu compute target...")

    gpu_target = AmlCompute(
        name=target_name,
        type="amlcompute",
        size="STANDARD_NC12S_V3",
        max_instances=5,
        idle_time_before_scale_down=300,
    )
    ml_client.compute.begin_create_or_update(gpu_target).result()

    print(
        f"AMLCompute with name {gpu_target.name} is created, the compute size is {gpu_target.size}"
    )

准备训练脚本

确保当前目录包含要在远程资源上运行的代码。这包括训练脚本及其所有依赖文件。在此示例中，提供了训练脚本

train_rapids.py- RAPIDS 环境的入口脚本，包括将数据集加载到 cuDF 数据框中，使用 Random Forest 进行训练，并使用 cuML 进行推理。

我们将在训练脚本中使用 mlflow 记录一些参数和指标，包括最高准确率

import mlflow

mlflow.log_metric('Accuracy', np.float(global_best_test_accuracy))

当我们开始在“调优模型超参数”部分中对模型进行超参数调优时，这些运行指标将变得尤为重要。

在远程计算上训练模型

设置环境

我们将使用自定义 RAPIDS Docker 镜像来设置环境。这可在 rapidsai/base 仓库的DockerHub 上获得。

%%bash
# create a Dockerfile defining the image the code will run in
cat > ./Dockerfile <<EOF
FROM nvcr.io/nvidia/rapidsai/base:25.04-cuda12.8-py3.12

RUN conda install --yes -c conda-forge 'dask-ml>=2024.4.4' \
 && pip install azureml-mlflow
EOF

确保您拥有正确的 Docker 构建上下文路径，即 os.getcwd()。

import os

from azure.ai.ml.entities import BuildContext, Environment

env_docker_image = Environment(
    build=BuildContext(path=os.getcwd()),
    name="rapids-hpo",
    description="RAPIDS environment with azureml-mlflow",
)

ml_client.environments.create_or_update(env_docker_image)

提交训练作业

我们将使用command类配置并运行训练作业。command 可用于运行独立作业或作为管道内部的函数。inputs 是一个字典，包含要传递给训练脚本的命令行参数。

from azure.ai.ml import Input, command

command_job = command(
    environment=f"{env_docker_image.name}:{env_docker_image.version}",
    experiment_name="test_rapids_aml_hpo_cluster",
    code=os.getcwd(),
    inputs={
        "data_dir": Input(type="uri_file", path=data_uri),
        "n_bins": 32,
        "compute": "single-GPU",  # multi-GPU for algorithms via Dask
        "cv_folds": 5,
        "n_estimators": 100,
        "max_depth": 6,
        "max_features": 0.3,
    },
    command="python train_rapids.py \
                    --data_dir ${{inputs.data_dir}} \
                    --n_bins ${{inputs.n_bins}} \
                    --compute ${{inputs.compute}} \
                    --cv_folds ${{inputs.cv_folds}} \
                    --n_estimators ${{inputs.n_estimators}} \
                    --max_depth ${{inputs.max_depth}} \
                    --max_features ${{inputs.max_features}}",
    compute=gpu_target.name,
)


# submit the command
returned_job = ml_client.jobs.create_or_update(command_job)

# get a URL for the status of the job
returned_job.studio_url

调优模型超参数

我们可以使用 Azure 机器学习的超参数调优功能来优化模型的超参数并提高准确率。

启动超参数扫描

让我们定义要扫描的超参数空间。我们将调优 n_estimators、max_depth 和 max_features 参数。在此示例中，我们将使用随机抽样来尝试不同的超参数配置集并最大化 Accuracy。

from azure.ai.ml.sweep import Choice, Uniform

command_job_for_sweep = command_job(
    n_estimators=Choice(values=range(50, 500)),
    max_depth=Choice(values=range(5, 19)),
    max_features=Uniform(min_value=0.2, max_value=1.0),
)

# apply sweep parameter to obtain the sweep_job
sweep_job = command_job_for_sweep.sweep(
    compute=gpu_target.name,
    sampling_algorithm="random",
    primary_metric="Accuracy",
    goal="Maximize",
)


# Relax these limits to run more trials
sweep_job.set_limits(
    max_total_trials=5, max_concurrent_trials=5, timeout=18000, trial_timeout=3600
)

# Specify your experiment details
sweep_job.display_name = "RF-rapids-sweep-job"
sweep_job.description = "Run RAPIDS hyperparameter sweep job"

这将启动 RAPIDS 训练脚本，并使用上面单元格中指定的参数。

# submit the hpo job
returned_sweep_job = ml_client.create_or_update(sweep_job)

监控运行

print(f"Monitor your job at {returned_sweep_job.studio_url}")

查找并注册最佳模型

下载最佳试验模型输出

ml_client.jobs.download(returned_sweep_job.name, output_name="model")

删除集群

ml_client.compute.begin_delete(gpu_target.name).wait()