如何防止附加到列表以填充 gpu 内存

Question

我运行一个带有批量推理的 pytorch 模型。我想将该模型的输出保存到列表中，因为稍后我需要它进行其他计算。问题是，如果我将模型的输出附加到我的列表“输出”，它会填满 gpu 内存。模型的输出是一个列表。有没有办法将列表移动到 RAM？

if __name__ == "__main__":
        output = []
        with torch.no_grad():
            for i in input_split:
                try:
                    preds = model(i)                        
                    output.append(preds)          
                    del preds   
                    gc.collect()
                    torch.cuda.empty_cache()    

                except RuntimeError as e:
                    print("Failed")

Answer 1

这是因为你从preds = model(i)得到的张量还在GPU中。

您可以在将它们附加到列表之前将它们从 GPU 中取出

output = []
with torch.no_grad():
    for i in input_split:
        preds = model(i)
        output.append(preds.cpu())

当你想在 GPU 中再次使用它们时，只需将它们一个一个地放入 GPU

for data in output:
    data = data.cuda()

编辑 1. Detectron2

这与上面的答案类似，只是有点复杂。

Detectron2的输出是list[dict]。这意味着如果批处理中有超过 1 张图像，则必须创建一个空列表作为缓冲区，例如

output = []
with torch.no_grad():
    for i in input_split:
        preds = model(i)
        
        buffer = []
        for pred in preds:
            pred['instances']['pred_boxes'] = pred['instances']['pred_boxes'].cpu()
            pred['instances']['scores'] = pred['instances']['scores'].cpu()
            pred['instances']['pred_classes'] = pred['instances']['pred_classes'].cpu()
            pred['instances']['pred_masks'] = pred['instances']['pred_masks'].cpu()
            pred['instances']['pred_keypoints'] = pred['instances']['pred_keypoints'].cpu()
            buffer.append(pred)
        output.append(buffer)

您或许也可以像这样使用指针并摆脱缓冲区

output = []
with torch.no_grad():
    for i in input_split:
        preds = model(i)
        
        for pred in preds:
            pred['instances']['pred_boxes'] = pred['instances']['pred_boxes'].cpu()
            pred['instances']['scores'] = pred['instances']['scores'].cpu()
            pred['instances']['pred_classes'] = pred['instances']['pred_classes'].cpu()
            pred['instances']['pred_masks'] = pred['instances']['pred_masks'].cpu()
            pred['instances']['pred_keypoints'] = pred['instances']['pred_keypoints'].cpu()
        output.append(preds)

如何防止附加到列表以填充 gpu 内存

How to prevent appending to a list to fill gpu memory

python

memory

gpu

pytorch

编辑 1. Detectron2