如何释放所有PyTorch占用的GPU内存？

Question

我有某种高级代码，因此模型训练等都被包装在pipeline_network类中。我主要的目标是每个新的折叠都训练一个新模型。

for train_idx, valid_idx in cv.split(meta_train[DEPTH_COLUMN].values.reshape(-1)):

        meta_train_split, meta_valid_split = meta_train.iloc[train_idx], meta_train.iloc[valid_idx]

        pipeline_network = unet(config=CONFIG, suffix = 'fold' + str(fold), train_mode=True)

但是当我进行第二次折叠时，所有内容都因为GPU内存不足而失败：

RuntimeError: cuda runtime error (2) : out of memory at /pytorch/torch/lib/THC/generic/THCStorage.cu:58

在时代结束时，我尝试手动删除该管道，但没有成功：

Translated:

在时代结束时，我尝试手动删除那个管道，但没有成功：

 def clean_object_from_memory(obj): #definition
    del obj
    gc.collect()
    torch.cuda.empty_cache()

clean_object_from_memory( clean_object_from_memory) # calling

调用这个也没有帮助：

def dump_tensors(gpu_only=True):
        torch.cuda.empty_cache()
        total_size = 0
        for obj in gc.get_objects():
            try:
                if torch.is_tensor(obj):
                    if not gpu_only or obj.is_cuda:
                        del obj
                        gc.collect()
                elif hasattr(obj, "data") and torch.is_tensor(obj.data):
                    if not gpu_only or obj.is_cuda:
                        del obj
                        gc.collect()
            except Exception as e:
                pass

我如何在进入下一个折叠之前重置PyTorch？

- Rocketq

Answer 1

尝试使用del删除对象，然后应用torch.cuda.empty_cache()。此操作后可释放可重复使用的内存。

- HzCheng

---

我也建议这一步骤。但你说得对，这是主要的步骤。 - Rocketq