Pytorch：GPU 内存泄漏

Question

我推测我在使用 PyTorch 框架训练 Conv 网络时遇到了 GPU 内存泄漏。下图

为了解决这个问题，我添加了 -

os.environ['CUDA_LAUNCH_BLOCKING'] = "1"

这样就解决了内存问题，如下图——

但由于我当时使用的是torch.nn.DataParallel，所以我希望我的代码能够利用所有 GPU，但现在它只利用GPU:1.

在使用之前os.environ['CUDA_LAUNCH_BLOCKING'] = "1"，GPU 利用率低于（同样糟糕）-

经过进一步挖掘，我发现，当我们使用 时torch.nn.DataParallel，我们不应该使用CUDA_LAUNCH_BLOCKING'，因为它会使网络陷入某种死锁机制。所以，现在我又回到了 GPU 内存问题，因为我认为我的代码没有利用它在没有设置的情况下显示的那么多内存CUDA_LAUNCH_BLOCKING=1。
我要使用的代码torch.nn.DataParallel-

device =  torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
if torch.cuda.device_count() > 1:
  print("Let's use", torch.cuda.device_count(), "GPUs!")
  # dim = 0 [30, xxx] -> [10, ...], [10, ...], [10, ...] on 3 GPUs
  model_transfer = nn.DataParallel(model_transfer.cuda(),device_ids=range(torch.cuda.device_count()))


model_transfer.to(device)

如何解决GPU内存问题？编辑：最少的代码 -

image_dataset = datasets.ImageFolder(train_dir_path,transform = transform)
train_loader = torch.utils.data.DataLoader(image_dataset['train'], batch_size=batch_size,shuffle = True)
model_transfer = models.resnet18(pretrained=True)

device =  torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
if torch.cuda.device_count() > 1:
  print("Let's use", torch.cuda.device_count(), "GPUs!")
  # dim = 0 [30, xxx] -> [10, ...], [10, ...], [10, ...] on 3 GPUs
  model_transfer = nn.DataParallel(model_transfer.cuda(),device_ids=range(torch.cuda.device_count()))
model_transfer.to(device) 

## Training function
for epoch in range(1, n_epochs+1):
    for batch_idx, (data, target) in enumerate(train_loader):
                    if use_cuda:
            data, target = data.to('cuda',non_blocking = True), target.to('cuda',non_blocking = True)
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output,target)
        loss.backward()
        optimizer.step()
        train_loss += ((1 / (batch_idx + 1)) * (loss.item() - train_loss))

   ## Validation loop same as training loop so not mentioning here

    print('Epoch: {} \tTraining Loss: {:.6f} \tValidation Loss: {:.6f}'.format(
        epoch,
        train_loss,
        valid_loss
        ))
        if valid_loss <= valid_loss_min:
           valid_loss_min,valid_loss))
           torch.save(model.state_dict(), 'case_3_model.pt')
           valid_loss_min = valid_loss

Answer 1

因此，我解决一些 CUDA 内存不足问题的方法是确保删除无用的张量并修剪可能因某些隐藏原因而被引用的张量。问题可能是由于请求的内存超出了您的容量，或者积累了不需要但不知何故留在内存中的垃圾数据。

内存管理最重要的方面之一是如何加载数据。与读取整个数据集不同，从磁盘读取（读取 npy 时使用 memmap）或进行批量加载（一次只读取一批图像或任何数据）可能会更有效地提高内存效率。尽管这可能会导致计算速度变慢，但它确实为您提供了灵活性，让您不必为了运行代码而外出购买更多 GPU 来存储内存。

我们不确定您的代码在读取数据或训练 CNN 方面的结构如何，所以这就是我能提供的建议。