使用 Dask 加速 GPU 集群上的 XGBoost

目录

在 GPU 集群上用 Dask 加载数据

提前停止训练

定制目标和评估指标

解释模型

运行推理

把它们放在一起

Scikit 学习包装器

---

在 xgboost1 . 0 中，我们引入了 新的官方 Dask 接口 来支持高效的分布式训练。 快速转发到 XGBoost1 . 4 ，接口现在功能齐全。如果您对 xgboostdask 接口还不熟悉，请参阅第一篇文章，以获得一个温和的介绍。在本文中，我们将看一些简单的代码示例，展示如何最大化 GPU 加速的好处。

我们的例子集中在希格斯数据集上，这是一个来自 机器学习库 的中等规模的分类问题。 在下面的章节中，我们从基本数据加载和预处理开始，使用 GPU 加速的 Dask 和 Dask-ml 。然后，针对不同配置的返回数据训练 XGBoost 模型。同时，分享一些新特性。之后，我们将展示如何在 GPU 集群上计算 SHAP 值以及可以获得的加速比。最后，我们分享了一些优化技术与推理。

以下示例需要在至少有一个 NVIDIA GPU 的机器上运行， GPU 可以是笔记本电脑或云实例。 Dask 的优点之一是它的灵活性，用户可以在笔记本电脑上测试他们的代码。它们还可以将计算扩展到具有最小代码更改量的集群。 另外，要设置环境，我们需要 xgboost==1.4 、 dask 、 dask-ml 、 dask-cuda 和 达斯克 – cuDF python 包，可从 RAPIDS 康达频道： 获得

conda install -c rapidsai -c conda-forge dask dask-ml dask-cuda dask-cudf xgboost=1.4.2

在 GPU 集群上用 Dask 加载数据

首先，我们将数据集下载到 data 目录中。

mkdir data
curl http://archive.ics.uci.edu/ml/machine-learning-databases/00280/HIGGS.csv.gz --output ./data/HIGGS.csv.gz

然后使用 dask-cuda 设置 GPU 集群：

import os
from time import time
from typing import Tuple

from dask import dataframe as dd
from dask_cuda import LocalCUDACluster
from distributed import Client, wait
import dask_cudf
from dask_ml.model_selection import train_test_split

import xgboost as xgb
from xgboost import dask as dxgb
import numpy as np
import argparse

# … main content to be inserted here in the following sections

if __name__ == "__main__":
   parser = argparse.ArgumentParser()
   parser.add_argument("--n_workers", type=int, required=True)
   args = parser.parse_args()
   with LocalCUDACluster(args.n_workers) as cluster:
       print("dashboard:", cluster.dashboard_link)
       with Client(cluster) as client:
           main(client)

给定一个集群，我们开始将数据加载到 gpu 中。 由于在参数调整期间多次加载数据，因此我们将 CSV 文件转换为 Parquet 格式以获得更好的性能。 这可以使用 dask_cudf 轻松完成：

def to_parquet() -> str:
   """Convert the HIGGS.csv file to parquet files."""
   dirpath = "./data"
   parquet_path = os.path.join(dirpath, "HIGGS.parquet")
   if os.path.exists(parquet_path):
       return parquet_path
   csv_path = os.path.join(dirpath, "HIGGS.csv")
   colnames = ["label"] + ["feature-%02d" % i for i in range(1, 29)]
   df = dask_cudf.read_csv(csv_path, header=None, names=colnames, dtype=np.float32)
   df.to_parquet(parquet