使用 RAPIDS 和 Dask 进行 HPO 基准测试#

超参数优化 (HPO) 通过探索可能的配置空间来帮助找到模型的最佳版本。虽然通常是理想的，但这种搜索计算成本高昂且耗时。

在下面的 Notebook 演示中，我们比较了基准测试结果，以展示 GPU 如何相对于 CPU 加速 HPO 调优作业。

例如，在对 3 年期航空公司数据集进行 100 次 HPO 试验时，我们将 p3.8xlarge Tesla V100 GPU 与 c5.24xlarge CPU EC2 实例进行比较发现，XGBoost 的实际时间加速了 48 倍（0.71 小时 对比 34.6 小时），RandomForest 加速了 16 倍（3.86 小时 对比 63.2 小时）。

前言

您可以设置本地环境，但建议启动虚拟机服务（Azure、AWS、GCP 等）。

对于此 Notebook，我们将使用 Amazon Machine Image (AMI) 作为起点。

参阅文档

请按照 AWS 弹性计算云 (EC2) 中的说明启动带有 GPU、NVIDIA 驱动程序和 NVIDIA 容器运行时的 EC2 实例。

注意

配置实例时，请确保在 “Amazon Machine Image (AMI)” 下的 AMI 选择框中选择 深度学习 AMI GPU TensorFlow 或 PyTorch

实例运行并访问 Jupyter 后，保存此 Notebook 并运行各单元格。

访问文档 >>

Python ML 工作流程

为了使用 RAPIDS 容器，入口点逻辑应解析参数、加载、预处理和分割数据、构建和训练模型、对训练好的模型进行评分/评估，并输出表示给定超参数设置的最终评分。

让我们逐步查看 ML 工作流程的每个阶段

数据集

我们利用了 Airline 数据集，这是一个大型的美国国内航班日志公共跟踪器，我们提供各种大小（1 年、3 年和 10 年）以及 Parquet（压缩列存储）格式。使用此数据集的机器学习目标是预测航班到达目的地是否会延误超过 15 分钟。

我们在 us-east-1 和 us-west-2 的公共 S3 演示存储桶中托管演示数据集。为了优化性能，我们建议您在与 EC2 实例相同的区域访问 S3 存储桶，以减少网络延迟和数据传输成本。

对于此演示，我们使用 3_year 数据集，其中包含以下部分特征

  • 日期和距离 ( 年、月、距离 )

  • 航空公司/承运人 ( Flight_Number_Reporting_Airline )

  • 实际出发和到达时间 ( DepTime 和 ArrTime )

  • 计划时间与实际时间之差 ( ArrDelay 和 DepDelay )

  • 延误的二进制编码版本，即我们的目标变量 ( ArrDelay15 )

配置 AWS 凭据以访问 S3 存储

aws configure

从 S3 存储桶下载数据集到当前工作目录

aws s3 cp --recursive s3://sagemaker-rapids-hpo-us-west-2/3_year/ ./data/

算法

从 ML/算法的角度来看，我们提供 XGBoost 和 RandomForest。您可以随意切换这些算法选项，示例中的一切都将继续有效。

parser = argparse.ArgumentParser()
parser.add_argument(
    "--model-type", type=str, required=True, choices=["XGBoost", "RandomForest"]
)

我们还可以通过重新打乱训练集/测试集划分（即交叉验证折叠）来选择性地提高鲁棒性。典型值为 3 到 10 折。我们将使用

n_cv_folds = 5

Dask 集群

为了最大限度地提高效率，我们启动一个用于 CPU 的 Dask LocalCluster 或利用 GPU 进行分布式计算的 LocalCUDACluster。然后连接 Dask Client 以在集群上提交和管理计算。

然后，我们可以摄取数据，并使用 dask 将其“持久化”到内存中，如下所示

if args.mode == "gpu":
    cluster = LocalCUDACluster()
else: # mode == "cpu"
    cluster = LocalCluster(n_workers=os.cpu_count())

with Client(cluster) as client:
    dataset = ingest_data(mode=args.mode)
    client.persist(dataset)

搜索范围

运行 HPO 时最重要的选择之一是选择超参数搜索过程的边界。在此 Notebook 中，我们利用了 Optuna 的强大功能，它是一个广泛用于超参数优化的 Python 库。

以下是 Optuna 入门的快速步骤

  1. 定义目标函数，它代表模型训练和评估过程。它将超参数作为输入，并返回要优化的指标（例如，在本例中是准确性）。请参阅 hpo.py 中的 train_xgboost() 和 train_randomforest()

  1. 使用 Trial 对象的方法指定搜索空间，以定义超参数及其相应的值范围或分布。例如

"max_depth": trial.suggest_int("max_depth", 4, 8),
"max_features": trial.suggest_float("max_features", 0.1, 1.0),
"learning_rate": trial.suggest_float("learning_rate", 0.001, 0.1, log=True),
"min_samples_split": trial.suggest_int("min_samples_split", 2, 1000, log=True),

  1. 创建一个 Optuna study 对象来跟踪试验及其相应的超参数配置和评估指标。

study = optuna.create_study(
        sampler=RandomSampler(seed=args.seed), direction="maximize"
    )

  1. 选择一种优化算法来确定 Optuna 如何探索和利用搜索空间以找到最优配置。例如，RandomSampler 是 Optuna 库提供的一种算法，它从搜索空间中随机采样超参数配置。

  2. 通过在 study 对象上调用 Optuna 的 optimize() 函数来运行优化。您可以指定要运行的试验次数或并行作业数。

 study.optimize(lambda trial: train_xgboost(
                    trial, dataset=dataset, client=client, mode=args.mode
                ),
                n_trials=100,
                n_jobs=1,
            )

运行 HPO

让我们试试吧！

此处包含的示例文件 hpo.py 实现了上述模式。

首先，通过运行 nvidia-smi 确保您拥有正确的 CUDAtoolkit 版本。有关支持的 GPU 和驱动程序范围的详细信息，请参阅 RAPIDS 安装文档（链接）。

!nvidia-smi

执行基准测试可能是一个艰巨且耗时的过程，可能需要数天。通过使用 tmux 等工具，您可以保持终端会话处于活动状态，即使 SSH 连接中断，也能确保任务继续运行。

tmux

运行以下命令可在 Docker 容器中运行超参数优化。

如果您的本地还没有该镜像，首次运行此命令可能需要几分钟来拉取。之后，启动会非常快。

以下命令中的参数作用如下

  • --gpus all = 使系统上的所有 GPU 可用于容器中的进程

  • --env EXTRA_CONDA_PACKAGES = 安装 optuna 和 optuna-integration conda 包

    • 该镜像已经安装了所有 RAPIDS 库及其依赖项

  • -p 8787:8787 = 将主机端口 8787 转发到容器端口 8787

    • 导航到 `{public IP of box}:8787` 查看 Dask 仪表盘！

  • -v / -w = 将主机上的当前目录挂载到容器中

    • 这允许容器中的进程读取您之前下载到 ./data 目录的数据

    • 这也意味着在容器内部对这些文件所做的任何更改都将反映回主机上

将输出管道重定向到名为 xgboost_hpo_logs.txt 的文件很有帮助，因为它可以保留所有日志以供后续检查。

!docker run \
    --gpus all \
    --env EXTRA_CONDA_PACKAGES="optuna optuna-integration" \
    -p 8787:8787 \
    -v $(pwd):/home/rapids/xgboost-hpo-example \
    -w /home/rapids/xgboost-hpo-example \
    -it nvcr.io/nvidia/rapidsai/base:25.04-cuda12.8-py3.12 \
    /bin/bash -c "python ./hpo.py --model-type 'XGBoost' --target 'gpu'" \
> ./xgboost_hpo_logs.txt 2>&1

尝试一些修改

现在您已经运行了这个示例，尝试一些修改吧！

例如

  • 使用 --model-type "RandomForest" 查看随机森林模型与 XGBoost 的比较情况

  • 使用 --target "cpu" 来估计 GPU 加速训练带来的速度提升

  • 使用其他自定义设置修改 hpo.py 中的管道