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() 8 Gb之后，使用optimizer.step() 12 Gb之后，即所有可用的内存。

在第二个历时正向运行正常，但是在反向传播期间，我遇到RuntimeError：CUDA错误：内存不足。

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

UPD：看来我的问题与此问题https://discuss.pytorch.org/t/how-to-free-gpu-memory-and-delete-memory-allocated-variables/20856

类似

Answer 1

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

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

Pytorch卷积网络内存使用情况详细信息

1 个答案: