Pytorch卷积网络内存使用详情

Question

我正在尝试为一个非常大的输入(5*100,000,000)训练一个神经网络，它需要比预期更多的内存。下面是一些最小的示例：

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import time

class Net(nn.Module):

    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv1d(in_channels=5, out_channels=1, kernel_size=100000000, stride=10)

    def forward(self, x):
        x = self.conv1(x)
        x = torch.sigmoid(x)
        return x

model = Net().cuda()

optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion = torch.nn.BCELoss()

data = torch.normal(torch.zeros(1,5,100000000),torch.ones(1,5,100000000))
data = data.cuda()
label = torch.ones(1,1,1)
label = label.cuda()

for epoch in range(10):
    output = model(data)
    loss = criterion(output, label)
   
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    print("Epoch :", epoch)

输入是一些随机数据，它使用了大约2 2Gb，正如预期的那样(32位*5* 100,000,000= 1.86 2Gb).This变量没有梯度。该网络由单个卷积层组成，其中一个滤波器的大小与输入相同，因此它具有500M的权重，即另一个2 2Gb。在前向传递之后，另有2 2Gb被使用。在使用loss.backprop() 8 Gb之后，在使用optimizer.step() 12 Gb之后，这是所有可用内存。

在第二个时期，正向传递运行正常，但在反向传播期间，我得到RuntimeError: CUDA error: out of memory。

在纪元期间，GPU内存中到底保存了什么？为什么优化步骤完成后没有释放内存？在这种情况下，如何减少内存使用量？

更新:看起来我的问题和这个问题类似，https://discuss.pytorch.org/t/how-to-free-gpu-memory-and-delete-memory-allocated-variables/20856

UPD2:在这里得到了pytorch开发人员的答案https://github.com/pytorch/pytorch/issues/12651，但它只是说这不是pytorch，而是cuDNN问题。

Answer 1

由于要调用loss.backprop()，因此PyTorch必须计算渐变，这导致了大量的内存分配。如果你想删除渐变，在一个变量上调用.detach()。

要释放未使用的内存，如果您想深入了解细节，可以调用torch.cuda.empty_cache()，the CUDA semantics page可能是一个起点。