如何释放所有内存pytorch取自gpu内存

Posted

技术标签:

【中文标题】如何释放所有内存pytorch取自gpu内存【英文标题】:How to free up all memory pytorch is taken from gpu memory 【发布时间】:2019-02-11 19:06:49 【问题描述】:

我有一些高级代码,所以模型训练等被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

在 epoch 结束时,我尝试手动删除该管道,但没有成功:

 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 然后我继续下一个折叠?

【问题讨论】:

【参考方案1】:

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

【讨论】:

我也建议了这一步。但你说得对,这是主要步骤

以上是关于如何释放所有内存pytorch取自gpu内存的主要内容,如果未能解决你的问题,请参考以下文章

pytorch 从 gpu 中删除模型

使用 pytorch 获取可用 GPU 内存的总量

如何避免 PyTorch 中的“CUDA 内存不足”

释放 Pycuda 的内存

需要OpenGL GPU内存清理吗?

Pytorch 模型在 CPU 和 GPU 上都内存不足,无法弄清楚我做错了啥