PyTorch 经常问的问题(FAQ)详解

2025-06-24 18:29 更新

一、模型报告“CUDA 运行时错误(2):内存不足”

当您的模型报告“CUDA 运行时错误(2):内存不足”时,这表明您的 GPU 内存已用完。以下是一些常见的检查事项:

1.1 不要在训练循环中累积历史记录

默认情况下,涉及需要渐变的变量的计算将保留历史记录。您应避免在计算中使用此类变量,这些变量将不受训练循环的影响,例如在跟踪统计信息时。相反,您应该分离变量或访问其基础数据。

示例代码

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 += float(loss)  # 解决方法

1.2 不要使用不需要的张量和变量

如果将 Tensor 或 Variable 分配给本地,Python 将不会取消分配,直到本地超出范围。您可以通过 del 释放此引用。同样,如果将 Tensor 或 Variable 分配给对象的成员变量,则在对象超出范围之前它不会释放。如果不需要的临时存储,将获得最佳的内存使用率。

示例代码

for i in range(5):
    intermediate = f(input[i])
    result += g(intermediate)
    del intermediate  # 释放引用
output = h(result)
return output

1.3 不要对太大的序列运行 RNN

通过 RNN 反向传播所需的内存量与 RNN 输入的长度成线性比例;因此,如果您尝试向 RNN 输入过长的序列,则会耗尽内存。这种现象的技术术语是到时间的反向传播(BPTT),可以通过截断 BPTT 来解决。

## 使用截断 BPTT
def repackage_hidden(h):
    """Wraps hidden states in new Tensors, to detach them from their history."""
    if isinstance(h, torch.Tensor):
        return h.detach()
    else:
        return tuple(repackage_hidden(v) for v in h)

1.4 请勿使用太大的线性图层

线性层 nn.Linear(m, n) 使用的内存量与要素数量成正比关系。例如,nn.Linear(10000, 10000) 将使用 10000 * 10000 = 100,000,000 个浮点数,这大约是 400MB 的内存(假设每个浮点数是 4 字节)。如果您同时需要存储梯度,则内存使用量将翻倍。以这种方式穿透内存非常容易。

二、我的 GPU 内存未正确释放

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

如果即使在退出 Python 后仍没有释放 GPU 内存,则很可能某些 Python 子进程仍然存在。您可以通过 ps -elf | grep python 找到它们,然后使用 kill -9 [pid] 手动将其杀死。

## 手动清理 GPU 缓存
torch.cuda.empty_cache()

三、我的数据加载器工作人员返回相同的随机数

您可能会使用其他库在数据集中生成随机数。例如,当通过 fork 启动工作程序子流程时,NumPy 的 RNG 被复制。可以通过 worker_init_fn 选项在工人中正确设置随机种子。

import numpy as np
import torch


def worker_init_fn(worker_id):
    np.random.seed(np.random.get_state()[1][0] + worker_id)


dataloader = torch.utils.data.DataLoader(
    dataset,
    num_workers=4,
    worker_init_fn=worker_init_fn
)

四、我的经常性网络无法使用数据并行性

ModuleDataParalleldata_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):
    def __init__(self):
        super().__init__()
        self.my_lstm = nn.LSTM(input_size=10, hidden_size=20, batch_first=True)


    def forward(self, padded_input, input_lengths):
        total_length = padded_input.size(1)  # 获取最大序列长度
        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)

此外,当批处理尺寸为 1(即 batch_first=False)且数据并行时,需要格外小心。在这种情况下,pack_padded_sequence 的第一个参数的形状将为 [T x B x *],并且应沿昏暗 1 分散,而第二个参数 input_lengths 的形状将为 [B],并且应沿昏暗 0 分散。

## 批处理尺寸为 1 的情况
padded_input = torch.randn(10, 1, 10)  # [T x B x *]
input_lengths = torch.tensor([10])      # [B]


packed_input = pack_padded_sequence(padded_input, input_lengths, batch_first=False)
packed_output, _ = m.my_lstm(packed_input)
output, _ = pad_packed_sequence(packed_output, batch_first=False, total_length=padded_input.size(0))

五、常见问题解答(Q&A)

Q1:如何避免在 PyTorch 中出现 CUDA 内存不足的问题?

A1:可以通过以下几种方式避免 CUDA 内存不足的问题:

  • 不要在训练循环中累积历史记录,避免不必要的内存占用。
  • 及时释放不需要的张量和变量,使用 del 语句。
  • 对于 RNN,使用截断的反向传播(truncated BPTT)来减少内存占用。
  • 避免使用过大的线性层,减少内存需求。
  • 使用 torch.cuda.empty_cache() 手动清理 GPU 缓存。

Q2:如何解决数据加载器工作进程返回相同随机数的问题?

A2:可以通过设置 worker_init_fn 来解决数据加载器工作进程返回相同随机数的问题。在 worker_init_fn 中为每个工作进程设置不同的随机种子。

def worker_init_fn(worker_id):
    np.random.seed(np.random.get_state()[1][0] + worker_id)


dataloader = DataLoader(
    dataset,
    num_workers=4,
    worker_init_fn=worker_init_fn
)

Q3:如何在使用数据并行性时解决 RNN 的大小不匹配问题?

A3:可以通过使用 pad_packed_sequencetotal_length 参数来确保所有设备上的输出序列长度一致。在定义 RNN 模块时,确保在 forward 方法中传递 total_length 参数。

output, _ = pad_packed_sequence(packed_output, batch_first=True, total_length=total_length)

Q4:如何检查和清理未释放的 GPU 内存?

A4:可以通过以下步骤检查和清理未释放的 GPU 内存:

  1. 使用 nvidia-smi 查看 GPU 内存使用情况。
  2. 如果发现内存未正确释放,可以使用 torch.cuda.empty_cache() 手动清理缓存。
  3. 如果问题仍然存在,可以检查是否有子进程未正确退出,并使用 kill -9 [pid] 手动杀死这些进程。

Q5:如何优化 RNN 的内存使用?

A5:可以通过以下几种方式优化 RNN 的内存使用:

  • 使用截断的反向传播(truncated BPTT)减少内存占用。
  • 减少序列长度,避免过长的序列导致内存不足。
  • 使用更小的批量大小,减少单次计算的内存需求。
  • 优化模型结构,减少不必要的计算和存储。

六、优化技巧与最佳实践

6.1 内存管理最佳实践

  • 避免不必要的张量复制:尽量减少张量的复制操作,因为这会增加内存占用。
  • 使用 inplace 操作:某些操作(如 add_())会在原地修改张量,减少内存使用。
  • 及时释放张量:使用 del 语句及时释放不再需要的张量。
  • 使用 torch.cuda.empty_cache():在需要时手动清理 GPU 缓存,释放未使用的内存。

6.2 数据加载器优化

  • 设置 worker_init_fn:确保数据加载器工作进程正确初始化随机种子,避免返回相同的随机数。
  • 使用 num_workers:合理设置 num_workers 参数,提高数据加载效率。
  • 优化数据预处理:尽量减少数据预处理的复杂度,提高数据加载速度。

6.3 模型优化

  • 使用混合精度训练:通过使用混合精度训练(AMP)减少内存占用并加速计算。
  • 优化模型结构:减少模型的复杂度,避免使用过大的线性层和 RNN 层。
  • 使用模型量化:对模型进行量化,减少内存占用和计算量。

七、总结与展望

通过本文的介绍,我们总结了 PyTorch 中常见的内存管理和数据加载问题,并提供了解决方案和优化技巧。希望这些内容能够帮助您更好地理解和使用 PyTorch,提高模型的训练效率和性能。

关注编程狮(W3Cschool)平台,获取更多 PyTorch 相关的优质教程和资源。

以上内容是否对您有帮助:
在线笔记
App下载
App下载

扫描二维码

下载编程狮App

公众号
微信公众号

编程狮公众号