为什么使用 dask_cudf DataFrame 的 KNearestNeighbors 的 cuml predict() 方法需要这么长时间？

Question

我有一个大型数据集（大约 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)

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

predictions = model.predict(X_test)

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

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

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

Answer 1

可以在此处找到 KNN 回归器分布式版本的文档。

以下是获得最佳性能的一些规则：

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

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

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

希望它有帮助！