错误信息:
RuntimeError: CUDA out of memory. Tried to allocate....
解决方法:
- 减小batch size
- 在测试的时候,使用 torch.no_grad()
with torch.no_grad():
output = model(inputs)
- 释放缓存
可以使用 torch.cuda.empty_cache()
try:
output = model(inputs)
except RuntimeError as exception:
if "out of memory" in str(exception):
print("WARNING: out of memory")
if hasattr(torch.cuda, 'empty_cache'):
torch.cuda.empty_cache()
else:
raise exception
- 其他可能的方法
- 清理一下RAM和GPU的空间
- 如果在Dataloader里使用了num_workers, 可以设成num_works=1
- 如果没有显示已经使用多少内存的信息,可能是pytorch, cuda之类的版本不匹配
网友评论