常见问题解答#

创建日期：2018 年 2 月 15 日 | 最后更新日期：2021 年 8 月 5 日

我的模型报告“cuda runtime error(2): out of memory”（CUDA 运行时错误(2)：内存不足）#

正如错误消息所暗示的，您的 GPU 内存已用尽。由于我们在 PyTorch 中经常处理大量数据，一个小错误就可能迅速导致程序耗尽所有 GPU 内存；幸运的是，这些情况下的修复通常很简单。以下是一些常见的检查项：

不要在训练循环中累积历史记录。默认情况下，涉及需要梯度的变量的计算会保留历史记录。这意味着您应避免在超出训练循环的计算中使用此类变量，例如在跟踪统计信息时。相反，您应该分离变量或访问其底层数据。

有时，出现可微分变量的情况可能不明显。考虑以下训练循环（摘自来源）

total_loss = 0
for i in range(10000):
    optimizer.zero_grad()
    output = model(input)
    loss = criterion(output)
    loss.backward()
    optimizer.step()
    total_loss += loss

在这里，total_loss 会在整个训练循环中累积历史记录，因为 loss 是一个具有 autograd 历史记录的可微分变量。您可以通过编写 total_loss += float(loss) 来修复此问题。

此问题的其他实例：1。

不要保留不需要的张量和变量。如果您将 Tensor 或 Variable 赋给一个局部变量，Python 在该局部变量超出作用域之前不会将其释放。您可以通过使用 del x 来释放此引用。同样，如果您将 Tensor 或 Variable 赋给对象的成员变量，它在对象超出作用域之前不会被释放。如果您不保留不需要的临时变量，将获得最佳的内存使用情况。

局部变量的作用域可能比您预期的要大。例如

for i in range(5):
    intermediate = f(input[i])
    result += g(intermediate)
output = h(result)
return output

在这里，即使在 h 执行期间，intermediate 仍然保持活动状态，因为它的作用域超出了循环的末尾。要更早地释放它，完成使用后应使用 del intermediate。

避免对过大的序列运行 RNN。反向传播通过 RNN 所需的内存量与 RNN 输入的长度成线性关系；因此，如果您尝试向 RNN 输入过长的序列，将会耗尽内存。

这种现象的技术术语是随时间反向传播 (BPTT)，有许多参考资料说明如何实现截断 BPTT，包括在词语言模型示例中；截断由此论坛帖子中描述的 repackage 函数处理。

不要使用过大的线性层。一个线性层 nn.Linear(m, n) 使用 $O(nm)$ 内存：也就是说，权重的内存需求与特征数量的平方成正比。通过这种方式耗尽内存非常容易（并且请记住，您需要的内存至少是权重大小的两倍，因为您还需要存储梯度。）

考虑使用检查点 (Checkpointing)。您可以使用检查点来权衡内存和计算。

我的 GPU 内存没有被正确释放#

PyTorch 使用缓存内存分配器来加速内存分配。因此，nvidia-smi 中显示的值通常不反映真实的内存使用情况。有关 GPU 内存管理的更多详细信息，请参阅内存管理。

如果即使在 Python 退出后 GPU 内存仍未释放，很可能是某些 Python 子进程仍然存活。您可以通过 ps -elf | grep python 找到它们，并使用 kill -9 [pid] 手动终止它们。

我的内存不足异常处理程序无法分配内存#

您的代码中可能有一些尝试从内存不足错误中恢复的逻辑。

try:
    run_model(batch_size)
except RuntimeError: # Out of memory
    for _ in range(batch_size):
        run_model(1)

但是发现当确实发生内存不足时，您的恢复代码也无法分配内存。这是因为 python 异常对象保持着对引发错误的堆栈帧的引用。这阻止了原始张量对象被释放。解决方案是将您的 OOM 恢复代码移到 except 子句之外。

oom = False
try:
    run_model(batch_size)
except RuntimeError: # Out of memory
    oom = True

if oom:
    for _ in range(batch_size):
        run_model(1)

我的数据加载器工作进程返回相同的随机数#

您很可能在数据集中使用其他库生成随机数，并且工作进程是通过 fork 启动的。请参阅 torch.utils.data.DataLoader 的文档，了解如何使用其 worker_init_fn 选项正确设置工作进程中的随机种子。

我的循环神经网络与数据并行性不兼容#

在使用 Module 和 DataParallel 或 data_parallel() 时，使用 pack sequence -> recurrent network -> unpack sequence 模式存在一个细微差别。输入到每个设备上 forward() 的内容将只是整个输入的一部分。由于 torch.nn.utils.rnn.pad_packed_sequence() 默认只填充到它所看到的（即特定设备上最长）的输入，因此在收集结果时会发生尺寸不匹配。因此，您可以利用 pad_packed_sequence() 的 total_length 参数来确保 forward() 调用返回相同长度的序列。例如，您可以编写

from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence

class MyModule(nn.Module):
    # ... __init__, other methods, etc.

    # padded_input is of shape [B x T x *] (batch_first mode) and contains
    # the sequences sorted by lengths
    #   B is the batch size
    #   T is max sequence length
    def forward(self, padded_input, input_lengths):
        total_length = padded_input.size(1)  # get the max sequence length
        packed_input = pack_padded_sequence(padded_input, input_lengths,
                                            batch_first=True)
        packed_output, _ = self.my_lstm(packed_input)
        output, _ = pad_packed_sequence(packed_output, batch_first=True,
                                        total_length=total_length)
        return output


m = MyModule().cuda()
dp_m = nn.DataParallel(m)

此外，当批次维度是 dim 1（即 batch_first=False）并使用数据并行时，需要额外小心。在这种情况下，pack_padded_sequence 的第一个参数 padding_input 的形状将是 [T x B x *]，应沿 dim 1 分散，但第二个参数 input_lengths 的形状将是 [B]，应沿 dim 0 分散。需要额外的代码来处理张量形状。