我推测我在使用 PyTorch 框架训练 Conv 网络时遇到了 GPU 内存泄漏。如下图
为了解决这个问题,我添加了 -
os.environ['CUDA_LAUNCH_BLOCKING'] = "1"
但由于我使用的是 torch.nn.DataParallel
,所以我希望我的代码能够利用所有 GPU,但现在它只利用 GPU:1
。
在使用os.environ['CUDA_LAUNCH_BLOCKING'] = "1"
之前,GPU利用率低于(同样糟糕)-
关于挖掘further ,我开始知道,当我们使用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
最佳答案
因此,我解决一些 CUDA 内存不足问题的方法是确保删除无用的张量并修剪可能因某些隐藏原因而被引用的张量。问题可能是由于请求的内存超出了您的容量,或者积累了不需要但不知何故留在内存中的垃圾数据。
内存管理最重要的方面之一是如何加载数据。与读取整个数据集不同,从磁盘读取(读取 npy 时使用 memmap)或进行批量加载(一次只读取一批图像或任何数据)可能会更有效地提高内存效率。尽管这可能会导致计算速度变慢,但它确实为您提供了灵活性,让您不必为了运行代码而外出购买更多 GPU 来存储内存。
我们不确定您的代码在读取数据或训练 CNN 方面的结构如何,所以这就是我能提供的建议。
关于deep-learning - torch : GPU Memory Leak,我们在Stack Overflow上找到一个类似的问题: https://stackoverflow.com/questions/61991467/