为什么在dask_cudf DataFrame中，KNearestNeighbors的cuml predict()方法花了这么长时间？

Question

我有一个很大的数据集(大约8000万行)，并且我正在使用带有dask_cudf DataFrame的cuml来训练KNearestNeighbors回归模型。

我使用了4个GPU，每个GPU的rmm_pool_size都是15 am：

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

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

Answer 1

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

以下是获得最佳性能所需遵循的一些规则：

1. 索引(X_train和y_train)应由以平衡方式分布在工作进程上的大型分区组成。

• query (X_test)最好由一些分区组成，这些分区的样本数量是batch_size参数的倍数。他们对工人的态度并不重要。

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

希望对您有所帮助！