Dask 和 cudf 循环内存不足错误

编程问答 2022-04-25

问题描述

我正在使用 dask 和 Rapidsai 在大型 (6.9GB) 数据集上运行 xgboost 模型。硬件是 4 个 2080 TI，每个具有 11 GB 内存。原始数据集有几十个已进行单热编码的目标列，因此我尝试运行一个循环，一次保留一个目标列，丢弃其余列，运行模型，然后重复。

这里有一些伪代码来展示这个过程：

with LocalCUDACluster(n_workers=4) as cluster:
     with Client(cluster) as client:
          raw_data = dask_cudf.read_csv(MyFile)
          d_train,d_test = 'function to drop all target columns but one,then create 
                            daskDMatrix objects for training and testing'
          model = 'Custom gridsearch function that loops through xgb.dask.train function and 
                  returns a dict of optimal params'

好的，所以如果我在任何单个目标列上运行它，它都可以正常工作 - 向上推到可用内存，但不会中断。如果我尝试在集群/客户端分配后循环执行此操作：

targets = ['target1','target2',...]
with LocalCUDACluster(n_workers=4) as cluster:
     with Client(cluster) as client:
          for target in targets:
               same code as above

它在循环第三次运行时中断，内存不足。相反，如果我在循环中分配集群/客户端：

targets = ['target1',...]
for target in targets:
     with LocalCUDACluster(n_workers=4) as cluster:
          with Client(cluster) as client:
             same code as above

效果很好。

所以，我猜我在循环运行中创建的一些对象显然是持久化的，对吧？因此，我尝试在每个循环结束时将 d_train 和 d_test 设置为空列表以清除内存，但这不起作用。我是否在循环之前或循环内部读入原始数据或任何其他更改也无关紧要。

似乎唯一可行的是在每个循环结束时重新初始化集群和客户端。

这是为什么？有没有更有效的方法来处理 dask 分区的循环中的内存？无论如何，我已经读过 dask 在循环中的效率非常低，所以也许我什至不应该这样做。

对使用 dask 或 Rapids 进行内存管理的任何见解都将不胜感激。

解决方法

作为第一步，您可以选择要读入 dask_cudf 的数据列。这可能会帮助您在迭代和创建时解决内存问题

targets = ['target1','target2',...]
with LocalCUDACluster(n_workers=4) as cluster:
     with Client(cluster) as client:
          for target in targets:
              dask_cudf.read_csv(MyFile,usecols=target)
              ... ## keep going with your ETL

这里是参考文档，如果你想。 cudf 和 dask_cudf 之间的 API 类似，就像 dask https://docs.rapids.ai/api/cudf/stable/api.html#cudf.io.csv.read_csv

专业提示：如果列数据足够小，您可能可以使用 1 个 GPU 并只使用 cuDF，因为调度程序的开销较少，因此性能会更高。

您可以做的另一件事是腌制每个 xgboost 模型，将其保存到磁盘，然后删除/覆盖 xgboost 模型，因为这些模型会占用您的 GPU 空间。然后，您可以在完成 ETL 的这一部分后重新阅读它。

如果您对超参数优化 (HPO) 感兴趣，you should check out our HPO notebooks and utilities，例如 https://github.com/rapidsai/cloud-ml-examples/blob/main/dask/kubernetes/Dask_cuML_Exploration.ipynb

此外，您可能不需要所有 with 语句，但我不知道您在此代码段之外在做什么，因此为了方便回答，我将它们留在 :) 中。>

cudfdaskdask-mlpythonrapids