如何在 PyTorch 中清除 Cuda 内存

Question

我正在尝试获取我已经训练过的神经网络的输出。输入是大小为 300x300 的图像。我使用的批量大小为 1，但在成功获得 25 张图像的输出后，我仍然收到 CUDA error: out of memory 错误。

我在网上搜索了一些解决方案，发现了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 中。有什么方法可以实现吗？

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

Answer 1

我知道我哪里出错了。我将发布解决方案作为其他可能遇到相同问题的人的答案。

基本上，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 空间。