解决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中。

相关推荐
quintin202521 小时前
用AI重构HR Tech:绚星绚才,将HR专业能力转化为业务增长引擎
人工智能·重构
恒点虚拟仿真1 天前
智能电网变电站综合自动化虚拟仿真实验
人工智能·智能电网·虚拟仿真实验·电力虚拟仿真·智能电网虚拟仿真
悠闲蜗牛�1 天前
云智融合:人工智能与云计算融合实践指南
人工智能·云计算
盼小辉丶1 天前
Wasserstein GAN(WGAN)
人工智能·神经网络·生成对抗网络
EasyCVR1 天前
视频融合平台EasyCVR在智慧水利中的实战应用:构建全域感知与智能预警平台
人工智能·音视频
DisonTangor1 天前
阿里开源Qwen3-Omni-30B-A3B三剑客——Instruct、Thinking 和 Captioner
人工智能·语言模型·开源·aigc
独孤--蝴蝶1 天前
AI人工智能-机器学习-第一周(小白)
人工智能·机器学习
西柚小萌新1 天前
【深入浅出PyTorch】--上采样+下采样
人工智能·pytorch·python
丁学文武1 天前
大语言模型(LLM)是“预制菜”? 从应用到底层原理,在到中央厨房的深度解析
人工智能·语言模型·自然语言处理·大语言模型·大模型应用·预制菜
fie88891 天前
基于MATLAB的声呐图像特征提取与显示
开发语言·人工智能