解决网络训练验证过程中显存增加的原因

最新推荐文章于 2024-07-28 09:34:41 发布

原创

最新推荐文章于 2024-07-28 09:34:41 发布 · 6.3k 阅读

36 ·

CC 4.0 BY-SA版权

文章标签：

#深度学习 #pytorch

在进行深度学习网络训练时，遇到显存不断增长导致OOM的问题。通过分析发现，参数保存、eval模式使用不当及未及时清理缓存可能是原因。解决方案包括：在保存参数时使用.item()，将model.eval()放在torch.no_grad()之后，以及使用torch.cuda.empty_cache()清理无用显存。

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

最近在训练网络时发现网络训练了几个epoch之后就会出现OOM
一开始以为是内存不够，后来才发现是在网络训练过程中，显存会不断的增加。
针对以上的问题，查找资料总结了三种有用的方式

训练过程过程中，保存参数加.item()
原代码：

def train_one_epoch(
    model, criterion, train_dataloader, optimizer, epoch, clip_max_norm
):
    model.train()
    device = next(model.parameters()).device
    train_loss = 0
   
    for i, d in enumerate(train_dataloader):
        d = d.to(device)

        optimizer.zero_grad()
        out_net = model(d)
        loss = criterion(out_net, d, epoch)
        train_loss += loss

        loss.backward()
        if clip_max_norm > 0:
            torch