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)

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

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

在 epoch 期间，GPU 内存中究竟保存了什么？为什么优化步骤完成后不释放内存？在这种情况下如何减少内存使用？

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

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

Answer 1

由于您要调用loss.backprop()，PyTorch 必须计算梯度，这会导致大量内存分配。如果你想删除渐变，请.detach()调用变量。

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