Pytorch 卷积网络内存使用详情

Posted 2023-04-15

【中文标题】Pytorch 卷积网络内存使用详情

【英文标题】：Pytorch convolutional network memory usage details

【发布时间】：2019-03-16 15:51:30

【问题描述】：

我正在尝试为非常大的输入 (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() 8Gb 后，使用optimizer.step() 12 Gb 后，就是所有可用内存。

在第二个 epoch 前向传递运行正常，但在反向传播期间我得到 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：在这里得到了 pytorch 开发人员的回答 https://github.com/pytorch/pytorch/issues/12651，但它只是说这不是 pytorch，而是 cuDNN 问题。

【问题讨论】：

你不能用pytorch分配一个允许一次使用的特定百分比的内存吗？我问是因为我在 tensorflow 中做过这个。除了使用 torch.cuda 中的一些方法，比如 empty_cache() 我不太确定。

看起来 pytorch 无法限制 GPU 内存使用 ***.com/questions/49529372/…。我试过empty_cache，没有效果

【参考方案1】：

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

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

【讨论】：

empty_cache() 不起作用。据我了解，您建议在我不再需要它们的渐变时分离变量，对吗？这个我试过了，没有效果

处理一个非常相似的问题已经有一段时间了，但我不完全记得解决方案了。

老实说，这更多的是尝试和错误，而不是认真理解 pytorch 分配 cuda 内存的方式。不幸的是，我现在没有电脑来尝试一些东西。您可以根据情况分离或删除张量。在此操作之后调用 empty_cache() 应该释放内存。请问为什么要使用卷积？应用内核大小等于输入大小的卷积与应用例如nn.Linear（即全连接层），除了填充等副作用。

问题中的示例纯粹是人为设计的，旨在显示我正在处理的问题。我尝试了与卷积等效的线性层，但它在这样的输入下运行良好，即使增加输入大小也不会导致与卷积相同的行为，副作用是探索的好主意。我真的很想了解记忆的去向。输入是 2gb，权重是 2gb，梯度增加了 2gb，也就是 6gb，剩下的 6gb 有什么用，我们能以某种方式去掉它的某些部分吗？

我已经尝试了更好的测试，看起来是错误，我在这里报告了它github.com/pytorch/pytorch/issues/12651也许有人会解释这种行为