如何在PyTorch中清除Cuda内存

  • Z时代
  • 2024-01-10
  • 分类：问答

我正在尝试获取已经训练过的神经网络的输出。输入的图像尺寸为300x300。我正在使用1的批处理大小，但是CUDA error: out of

memory在成功获取25张图像的输出后，仍然出现错误。

我在网上搜索了一些解决方案，然后遇到了torch.cuda.empty_cache()。但这似乎仍无法解决问题。

这是我正在使用的代码。

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
train_x = torch.tensor(train_x, dtype=torch.float32).view(-1, 1, 300, 300)
train_x = train_x.to(device)
dataloader = torch.utils.data.DataLoader(train_x, batch_size=1, shuffle=False)
right = []
for i, left in enumerate(dataloader):
    print(i)
    temp = model(left).view(-1, 1, 300, 300)
    right.append(temp.to('cpu'))
    del temp
    torch.cuda.empty_cache()

这种for loop运行给内存错误之前为25次，每次时间。

每次，我都会在网络中发送一个新图像进行计算。因此，在循环中的每次迭代之后，我实际上都不需要将先前的计算结果存储在GPU中。有什么办法可以做到这一点？

任何帮助将不胜感激。谢谢。

回答：

我弄清楚哪里出了问题。我将解决方案发布为其他可能正在努力解决同一问题的人的答案。

基本上，PyTorch所做的是，每当我通过网络传递数据时，它都会创建一个计算图，并将计算结果存储在GPU内存中，以防万一我想在反向传播期间计算梯度。但是由于我只想执行前向传播，因此只需要torch.no_grad()为模型指定即可。

因此，我代码中的for循环可以重写为：

for i, left in enumerate(dataloader):
    print(i)
    with torch.no_grad():
        temp = model(left).view(-1, 1, 300, 300)
    right.append(temp.to('cpu'))
    del temp
    torch.cuda.empty_cache()

no_grad()为我的模型指定告诉PyTorch我不想存储任何以前的计算，从而释放了我的GPU空间。

以上是 如何在PyTorch中清除Cuda内存 的全部内容， 来源链接： utcz.com/qa/413850.html

回到顶部