PyTorch 经常问的问题(FAQ)详解
一、模型报告“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
)
四、我的经常性网络无法使用数据并行性
在 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):
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_sequence
的 total_length
参数来确保所有设备上的输出序列长度一致。在定义 RNN 模块时,确保在 forward
方法中传递 total_length
参数。
output, _ = pad_packed_sequence(packed_output, batch_first=True, total_length=total_length)
Q4:如何检查和清理未释放的 GPU 内存?
A4:可以通过以下步骤检查和清理未释放的 GPU 内存:
- 使用
nvidia-smi
查看 GPU 内存使用情况。 - 如果发现内存未正确释放,可以使用
torch.cuda.empty_cache()
手动清理缓存。 - 如果问题仍然存在,可以检查是否有子进程未正确退出,并使用
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 相关的优质教程和资源。
更多建议: