GTX1070 8GB下深度学习模型验证阶段CUDA显存不足问题求助

问题描述 #

我的深度学习模型在验证阶段出现CUDA error: out of memory错误，但训练阶段（包含前向/反向传播，本应因存储梯度占用更多显存）却未出现该问题。模型共21257650个参数，使用GTX 1070（8GB显存）进行训练。

错误信息 #

Traceback (most recent call last):
  File "main.py", line 303, in <module>
    main(args)
  File "main.py", line 284, in main
    val_loss, val_psnr = validation(args, epoch, writer)
  File "main.py", line 144, in validation
    for i, (images, gt_image) in loop:
  File "C:\Users\Anh\anaconda3\envs\Kienv1\lib\site-packages\tqdm\std.py", line 1182, in __iter__
    for obj in iterable:
  File "C:\Users\Anh\anaconda3\envs\Kienv1\lib\site-packages\torch\utils\data\dataloader.py", line 681, in __next__
    data = self._next_data()
  File "C:\Users\Anh\anaconda3\envs\Kienv1\lib\site-packages\torch\utils\data\dataloader.py", line 1376, in _next_data
    return self._process_data(data)
  File "C:\Users\Anh\anaconda3\envs\Kienv1\lib\site-packages\torch\utils\data\dataloader.py", line 1402, in _process_data      
    data.reraise()
  File "C:\Users\Anh\anaconda3\envs\Kienv1\lib\site-packages\torch\_utils.py", line 461, in reraise
    raise exception
RuntimeError: Caught RuntimeError in pin memory thread for device 0.
Original Traceback (most recent call last):
  File "C:\Users\Anh\anaconda3\envs\Kienv1\lib\site-packages\torch\utils\data\_utils\pin_memory.py", line 34, in _pin_memory_loop
    data = pin_memory(data, device)
  File "C:\Users\Anh\anaconda3\envs\Kienv1\lib\site-packages\torch\utils\data\_utils\pin_memory.py", line 65, in pin_memory    
    return type(data)([pin_memory(sample, device) for sample in data])  # type: ignore[call-arg]
  File "C:\Users\Anh\anaconda3\envs\Kienv1\lib\site-packages\torch\utils\data\_utils\pin_memory.py", line 65, in <listcomp>    
    return type(data)([pin_memory(sample, device) for sample in data])  # type: ignore[call-arg]
  File "C:\Users\Anh\anaconda3\envs\Kienv1\lib\site-packages\torch\utils\data\_utils\pin_memory.py", line 65, in pin_memory    
    return type(data)([pin_memory(sample, device) for sample in data])  # type: ignore[call-arg]
  File "C:\Users\Anh\anaconda3\envs\Kienv1\lib\site-packages\torch\utils\data\_utils\pin_memory.py", line 65, in <listcomp>    
    return type(data)([pin_memory(sample, device) for sample in data])  # type: ignore[call-arg]
  File "C:\Users\Anh\anaconda3\envs\Kienv1\lib\site-packages\torch\utils\data\_utils\pin_memory.py", line 50, in pin_memory    
    return data.pin_memory(device)
RuntimeError: CUDA error: out of memory
CUDA kernel errors might be asynchronously reported at some other API call,so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.

验证代码 #

def validation(args, epoch, writer):
    torch.cuda.empty_cache()
    losses, psnrs, ssims = utils.init_meters(args.loss, reset_loss=True)
    model.eval()
    criterion.eval()

    args.out_counter = 0  # 重置输出图像索引
    start = time.time()
    with torch.no_grad(): #no_grad()告知PyTorch无需存储过往计算记录，从而释放GPU空间
        loop = tqdm(enumerate(val_loader), total=len(val_loader))
        for i, (images, gt_image) in loop:
            # 构建输入批次
            images = [img_.to(device) for img_ in images]
            gt = [gt_img.to(device) for gt_img in gt_image]
            print("gt len: " + str(len(gt)))

            # 前向传播
            out = model(images)  # out = [framet1, framet2]
            if device.type == "cuda":
                print(f"GPU Memory Usage: {torch.cuda.memory_allocated() / 1024 ** 3:.2f} GB")
            loss = 0
            loss_specific = []

            # 计算每个输出的损失
            for output, target in zip(out, gt):
                single_loss, single_loss_specific = criterion(output, target)
                loss += single_loss.item()  # 使用loss.item()而非loss，避免占用梯度显存，曾解决过显存不足问题
                loss_specific.append(single_loss_specific)

            # 保存损失值
            for k, v in losses.items():
                if k != 'total':
                    v.update(loss_specific[k].item())
            losses['total'].update(loss.item())

            # 计算指标
            for output, target in zip(out, gt):
                utils.eval_metrics(output, target, psnrs, ssims)

            # Tensorboard记录
            if i % args.log_iter == 0:
                utils.log_tensorboard(writer, losses, psnrs.avg, ssims.avg.item(),
                                      optimizer.param_groups[0]['lr'], epoch * len(train_loader) + i, 'val')

            # 保存结果图像
            # if epoch % 15 == 0:
            #     epoch_path = os.path.join(args.save_path, args.result_images_folder, 'Epoch_' + str(epoch))
            #     utils.save_batch_images(args, out, gt, epoch_path)

            # 更新进度条
            loop.set_description("(Val)")
            loop.set_postfix(loss=loss.item())

        # 计算验证执行时间
        val_time_elapsed = time.time() - start

        # 将验证指标保存至csv
        if epoch % 2 == 0:
            utils.save_metrics(args, os.path.join(args.save_path, args.graph_folder), epoch,
                               losses['total'].avg, psnrs.avg, ssims.avg.item(),
                               optimizer.param_groups[0]['lr'], val_time_elapsed, 'val')

        print('验证结果: \t验证损失: {:.6f}\t验证耗时: {:.2f}'
              '\tPSNR: {:.4f}\tSSIM: {:.3f}'.format(losses['total'].avg, val_time_elapsed,
                                                    psnrs.avg, ssims.avg.item()))

        return losses['total'].avg, psnrs.avg

解决方法 #

• 降低验证集batch size：训练阶段可能用了梯度累积或较小batch size，验证阶段若batch size与训练一致，容易因显存未完全释放导致溢出。可将验证集batch size减半，比如训练用8则验证用4。
• 关闭pin_memory：错误栈显示问题出在pin_memory线程，验证阶段可将val_loader的pin_memory设为False，牺牲一点数据加载速度换取显存空间。
• 彻底清理训练阶段显存残留：在调用validation函数前，手动删除训练相关临时张量并清空缓存：

  del train_images, train_gt
  torch.cuda.empty_cache()
  

• 循环内及时释放张量：在验证循环末尾，手动删除当前批次的输出、标签等张量，避免显存累积：

  del out, gt, loss_specific
  torch.cuda.empty_cache()
  

• 检查工具函数显存占用：确认utils.eval_metrics和utils.log_tensorboard中未创建不必要的持久化张量，确保所有计算都在torch.no_grad()上下文内执行。
• 限制GPU显存分配比例：在代码开头添加以下代码，限制当前进程仅使用90%的GPU显存，避免与其他进程冲突：

  torch.cuda.set_per_process_memory_fraction(0.9, device=device)
  

内容的提问来源于stack exchange，提问作者AhNhat