如何在PyTorch中释放GPU内存
我使用以下代码使用几个模型，试图计算困惑的句子列表：
from transformers import AutoModelForMaskedLM, AutoTokenizer
import torch
import numpy as np
model_name = 'cointegrated/rubert-tiny'
model = AutoModelForMaskedLM.from_pretrained(model_name).cuda()
tokenizer = AutoTokenizer.from_pretrained(model_name)
def score(model, tokenizer, sentence):
tensor_input = tokenizer.encode(sentence, return_tensors='pt')
repeat_input = tensor_input.repeat(tensor_input.size(-1)-2, 1)
mask = torch.ones(tensor_input.size(-1) - 1).diag(1)[:-2]
masked_input = repeat_input.masked_fill(mask == 1, tokenizer.mask_token_id)
labels = repeat_input.masked_fill( masked_input != tokenizer.mask_token_id, -100)
with torch.inference_mode():
loss = model(masked_input.cuda(), labels=labels.cuda()).loss
return np.exp(loss.item())
print(score(sentence='London is the capital of Great Britain.', model=model, tokenizer=tokenizer))
# 4.541251105675365大多数模型运行良好，但有些句子似乎会造成错误：
RuntimeError: CUDA out of memory. Tried to allocate 10.34 GiB (GPU 0; 23.69 GiB total capacity; 10.97 GiB already allocated; 6.94 GiB free; 14.69 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation. See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF
这是有道理的，因为有些很长。所以我所做的就是添加一些类似try, except RuntimeError, pass的东西。
这似乎可以工作到大约210句话，然后它只输出错误：
CUDA error: an illegal memory access was encountered CUDA kernel errors might be asynchronously reported at some other API call,so the stacktrace below might be incorrect. For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
我发现这有很多的讨论和想法，有些是关于潜在的错误GPU？但我知道我的GPU的工作原理，因为这个精确的代码适用于其他型号。还有关于批处理大小的这里，这就是为什么我认为它可能与释放内存有关。
我试着运行torch.cuda.empty_cache()来释放内存，比如每隔一段时间就释放这里中的内存，但是它没有工作(抛出了相同的错误)。
更新：I过滤了长度超过550的句子，这似乎消除了CUDA error: an illegal memory access was encountered CUDA kernel errors might be asynchronously reported at some other API call,so the stacktrace below might be incorrect. For debugging consider passing CUDA_LAUNCH_BLOCKING=1.错误。
回答 2
Stack Overflow用户
发布于 2021-12-31 10:47:39
您需要在gc.collect()之前应用torch.cuda.empty_cache()，我也将模型拉到torch.cuda.empty_cache()上，然后删除该模型及其检查点。试试看什么对你有用：
import gc
model.cpu()
del model, checkpoint
gc.collect()
torch.cuda.empty_cache()Stack Overflow用户
发布于 2022-01-06 11:05:39
我没有确切的答案，但我可以分享我在类似的situations...hope中采用的一些故障排除技术，这可能会有所帮助。
- 首先，遗憾的是，有时
CUDA error是模糊的，所以您应该考虑在CPU上运行您的代码，看看是否真的有其他事情发生(参见这里 )。 - 如果问题与内存有关，下面是我使用的两个自定义实用程序：
from torch import cuda
def get_less_used_gpu(gpus=None, debug=False):
"""Inspect cached/reserved and allocated memory on specified gpus and return the id of the less used device"""
if gpus is None:
warn = 'Falling back to default: all gpus'
gpus = range(cuda.device_count())
elif isinstance(gpus, str):
gpus = [int(el) for el in gpus.split(',')]
# check gpus arg VS available gpus
sys_gpus = list(range(cuda.device_count()))
if len(gpus) > len(sys_gpus):
gpus = sys_gpus
warn = f'WARNING: Specified {len(gpus)} gpus, but only {cuda.device_count()} available. Falling back to default: all gpus.\nIDs:\t{list(gpus)}'
elif set(gpus).difference(sys_gpus):
# take correctly specified and add as much bad specifications as unused system gpus
available_gpus = set(gpus).intersection(sys_gpus)
unavailable_gpus = set(gpus).difference(sys_gpus)
unused_gpus = set(sys_gpus).difference(gpus)
gpus = list(available_gpus) + list(unused_gpus)[:len(unavailable_gpus)]
warn = f'GPU ids {unavailable_gpus} not available. Falling back to {len(gpus)} device(s).\nIDs:\t{list(gpus)}'
cur_allocated_mem = {}
cur_cached_mem = {}
max_allocated_mem = {}
max_cached_mem = {}
for i in gpus:
cur_allocated_mem[i] = cuda.memory_allocated(i)
cur_cached_mem[i] = cuda.memory_reserved(i)
max_allocated_mem[i] = cuda.max_memory_allocated(i)
max_cached_mem[i] = cuda.max_memory_reserved(i)
min_allocated = min(cur_allocated_mem, key=cur_allocated_mem.get)
if debug:
print(warn)
print('Current allocated memory:', {f'cuda:{k}': v for k, v in cur_allocated_mem.items()})
print('Current reserved memory:', {f'cuda:{k}': v for k, v in cur_cached_mem.items()})
print('Maximum allocated memory:', {f'cuda:{k}': v for k, v in max_allocated_mem.items()})
print('Maximum reserved memory:', {f'cuda:{k}': v for k, v in max_cached_mem.items()})
print('Suggested GPU:', min_allocated)
return min_allocated
def free_memory(to_delete: list, debug=False):
import gc
import inspect
calling_namespace = inspect.currentframe().f_back
if debug:
print('Before:')
get_less_used_gpu(debug=True)
for _var in to_delete:
calling_namespace.f_locals.pop(_var, None)
gc.collect()
cuda.empty_cache()
if debug:
print('After:')
get_less_used_gpu(debug=True)2.1 free_memory允许您将gc.collect和cuda.empty_cache组合起来，从命名空间中删除一些想要的对象，并释放它们的内存(您可以传递一个变量名列表作为to_delete参数)。这很有用，因为您可能有未使用的对象占用内存。例如，假设您遍历了3个模型，那么当您进入第二次迭代时，第一个模型可能仍然占用一些gpu内存(我不知道为什么，但我在笔记本中经历过这种情况，我能找到的唯一解决方案是重新启动笔记本或显式释放内存)。但是，我不得不说，这并不总是实用的，因为您需要知道哪些变量持有GPU memory...and，但情况并不总是如此，特别是当您有许多内部与模型相关的梯度时(更多信息请参见这里 )。您也可以尝试使用with torch.no_grad():而不是with torch.inference_mode():；它们应该是等价的，但可能值得一试.
2.2如果您有多个GPU环境，您可以考虑交替切换到使用较少的GPU，这要归功于其他实用程序get_less_used_gpu。
- 此外，您还可以尝试跟踪GPU的使用情况，以查看何时发生错误并从那里进行调试。如果您在linux平台上，我建议最好的/最简单的方法是使用nvtop。
希望一些有用的东西:)
https://stackoverflow.com/questions/70508960
复制相似问题
腾讯云开发者
