解决pytorch训练的过程中内存一直增加的问题

来自:解决pytorch训练的过程中内存一直增加的问题 - 知乎

pytorch训练中内存一直增加的原因(部分)

  • 代码中存在累加loss,但每步的loss没加item()

    import torch
    import torch.nn as nn
    from collections import defaultdict

    if torch.cuda.is_available():
    device = 'cuda'
    else:
    device = 'cpu'

    model = nn.Linear(100, 400).to(device)
    criterion = nn.L1Loss(reduction='mean').to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

    train_loss = defaultdict(float)
    eval_loss = defaultdict(float)

    for i in range(10000):
    model.train()
    x = torch.rand(50, 100, device=device)
    y_pred = model(x) # 50 * 400
    y_tgt = torch.rand(50, 400, device=device)

    复制代码
      loss = criterion(y_pred, y_tgt)
      optimizer.zero_grad()
      loss.backward()
      optimizer.step()
      
      # 会导致内存一直增加,需改为train_loss['loss'] += loss.item()
      train_loss['loss'] += loss
    
      if i % 100 == 0:
          train_loss = defaultdict(float)
          model.eval()
          x = torch.rand(50, 100, device=device)
          y_pred = model(x) # 50 * 400
    
          y_tgt = torch.rand(50, 400, device=device)
          loss = criterion(y_pred, y_tgt)
    
          # 会导致内存一直增加,需改为eval_loss['loss'] += loss.item()
          eval_loss['loss'] += loss

以上代码会导致内存占用越来越大,解决的方法是:train_l oss['loss'] += loss.item() 以及 eval_loss['loss'] += loss.item()。值得注意的是,要复现内存越来越大的问题,模型中需要切换model.train() 和 model.eval(),train_loss以及eval_loss的作用是保存模型的平均误差(这里是累积误差),保存到tensorboard中。

相关推荐
橡晟4 小时前
深度学习入门:让神经网络变得“深不可测“⚡(二)
人工智能·python·深度学习·机器学习·计算机视觉
墨尘游子4 小时前
神经网络的层与块
人工智能·python·深度学习·机器学习
Leah01054 小时前
什么是神经网络,常用的神经网络,如何训练一个神经网络
人工智能·深度学习·神经网络·ai
Leah01055 小时前
机器学习、深度学习、神经网络之间的关系
深度学习·神经网络·机器学习·ai
PyAIExplorer5 小时前
图像亮度调整的简单实现
人工智能·计算机视觉
Striker_Eureka6 小时前
DiffDet4SAR——首次将扩散模型用于SAR图像目标检测,来自2024 GRSL(ESI高被引1%论文)
人工智能·目标检测
Rvelamen6 小时前
LLM-SECURITY-PROMPTS大模型提示词攻击测评基准
人工智能·python·安全
AI technophile7 小时前
OpenCV计算机视觉实战(15)——霍夫变换详解
人工智能·opencv·计算机视觉
JNU freshman7 小时前
计算机视觉 之 数字图像处理基础(一)
人工智能·计算机视觉
鹧鸪云光伏8 小时前
鹧鸪云重构光伏发电量预测的精度标准
人工智能·无人机·光伏·光伏设计·光伏模拟