单 GPU 上的 Dask 推理 MMDetection

Question

我有一个预训练的 MMdetection 模型，我计划将其用于对大型数据集进行推理。我收集要在数据框中处理的数据集并将其传递到我的推理阶段。

# Loading model in gpu
det_model.cuda()

# Doing inference
def inference(batch_data):
    return det_model(batch_data)

with LocalCluster(n_worker=1, process=True, n_threads=1, memory='2GB') as cluster, Clinet(cluster) as client):
    data_prep = data_prep.map_partitions(inference, meta=dict, chunks=2)
    data_prep = data_prep.compute()

上述代码导致 OOM 错误。模型大小

det_model

~1GB。我已经尝试了多种方法来调整调度程序的

memory

限制，但没有任何帮助。

看一下 dask 上的例子Moving-the-model-around，我看到模型保留在CPU上。

我对深度学习模型很陌生，不太确定发生了什么。有人可以解释/指出正确的方向吗？？