PyTorch out of memory 解决方案

最新推荐文章于 2025-06-28 18:54:53 发布

bllddee

最新推荐文章于 2025-06-28 18:54:53 发布

阅读量5.7k

点赞数 4

CC 4.0 BY-SA版权

分类专栏： Python机器学习

本文链接：https://blog.youkuaiyun.com/Dooonald/article/details/89875620

Python机器学习专栏收录该内容

25 篇文章

订阅专栏

博客介绍了解决模型运行时Out of Memory问题的方法。一是当运行一段时间出现该问题，可能是无用临时变量过多，可用torch.cuda.empty_cache()清理；二是若一批数据中有长句易致OOM，可跳过该句子。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

如果模型在运行了一些时间后出现的outofmemory，那么有可能是因为无用的临时变量太多了，我们需要使用torch.cuda.empty_cache()进行清理就可以了。

try:
    loss, outputs = model(src, lengths, dec, targets)
except RuntimeError as e:
    if 'out of memory' in str(e):
        print('| WARNING: ran out of memory')
        if hasattr(torch.cuda, 'empty_cache'):
            torch.cuda.empty_cache()
    else:
        raise e

如果这一批数据中有一个很长的句子容易导致oom那么就跳过这个句子

except RuntimeError as e:
    if 'out of memory' in str(e):
        print('| WARNING: ran out of memory, skipping batch')
        ooms += 1
        self.zero_grad()
    else:
        raise e