为什么使用dask_cudf DataFrame的KNearestNeighbors的cuml预测方法花费这么长时间？

问题描述

我有一个大数据集（大约8000万行），并且正在使用dask_cudf DataFrame的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小时，终于一次看到了预测方法的结果。同样很明显，在运行预测方法期间，GPU利用率要低得多，下降到大约30-40％（在训练过程中约为100％），请参见以下屏幕截图：

我可以使用一些帮助来理解为什么预测方法要花这么长时间，以及我在代码中做错了什么。作为参考，我遵循此文档站点上提供的KNN Regressor示例：https://docs.rapids.ai/api/cuml/stable/api.html#id23

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

解决方法

可以在here中找到有关KNN Regressor分布式版本的文档。

要获得最佳性能，请遵循以下规则：

1. 索引（X_train和y_train）应该由在工作人员上平衡分配的大分区组成。

2. 查询（X_test）最好由分区组成，这些分区的样本数量是batch_size参数的倍数。他们对工人的处置并不重要。

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

希望它会有所帮助！