1

我有一个大型数据集（大约 8000 万行），我正在使用带有 dask_cudf 数据帧的 cuml 训练一个 KNearestNeighbors 回归模型。

我正在使用 4 个 GPU，每个 rmm_pool_size 为 15GB：

from dask.distributed import Client
from dask_cuda import LocalCUDACluster
import cudf, cuml
import dask_cudf

cluster = LocalCUDACluster(
    rmm_pool_size="15GB"
)

client = Client(cluster)
client.run(cudf.set_allocator, "managed")

我正在从存储在 S3 存储桶中的镶木地板文件中读取数据：

df = dask_cudf.read_parquet("s3://path-to-parquet/", chunksize="2 GB", dtype=dtypes)

当我拟合 KNN 模型时，它运行良好，我可以看到这段时间 GPU 利用率很高。这是我用来拟合模型的代码：

from cuml.dask.neighbors import KNeighborsRegressor
from dask_ml.model_selection import train_test_split    

target = "target_lat"
X = train_df.drop(columns=target)
y = train_df[target]

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)
print(X_train.shape)
print(y_train.shape)
print(X_test.shape)
print(y_test.shape)

model = KNeighborsRegressor(n_neighbors=5, client=client)
model.fit(X_train, y_train)

但是，当我尝试输出测试集的预测时，与 fit 方法相比，这需要大量时间。

predictions = model.predict(X_test)

我等了将近 24 小时才终于有一次看到 predict 方法的结果。也很明显，predict 方法运行期间的 GPU 利用率要低得多，下降到大约 30-40%（在训练期间约为 100%），见下面的截图：

预测方法期间的 CPU 和 GPU 利用率

我可以使用一些帮助来理解为什么 predict 方法需要这么长时间，以及我在代码中是否做错了什么。作为参考，我正在关注本文档站点上给出的 KNN Regressor 示例：https ://docs.rapids.ai/api/cuml/stable/api.html#id23

任何帮助将不胜感激，谢谢！

4

1 回答 1

1

KNN Regressor 的分布式版本的文档可以在这里找到。

为了获得最佳性能，请遵循以下几条规则：

  1. 索引（X_train 和 y_train）应该由以平衡方式分布在 worker 上的大分区组成。

  2. 查询 (X_test) 最好由具有多个样本的分区组成，这些样本数量是 batch_size 参数的倍数。他们对工人的态度并不重要。

  3. 设置一次处理多少个查询的 batch_size 参数可以设置为更高的值。

希望对您有所帮助！

于 2020-11-26T18:43:06.167 回答