torch.cuda.OutOfMemoryError 在 GPU 上训练模型时出现，但在 CPU 上训练较大批量时不会出现

Question

我正在努力在 PyTorch 中训练多模式模型。我有一个训练循环，它在我的 CPU 上运行得很好（尽管很慢）（我测试到批量大小 = 32）。然而，当我尝试在 GPU (Tesla P40) 上运行它时，它只能在批量大小 = 2 的情况下运行。如果批量大小较大，它会抛出 torch.cuda.OutOfMemoryError。我正在处理预嵌入的视频和音频以及预标记化的文本。 GPU 是否真的无法处理大于 2 的批量大小，或者我的代码是否有问题？您对我如何解决故障有什么建议吗？对于这个简单的问题我深表歉意，这是我第一次使用 GPU 集群。我正在大学的 GPU 集群上运行此代码，并仔细检查了我正在使用的 GPU 没有被其他人使用。

我尝试检查 CPU 和 GPU 上的内存使用情况。通过运行nvidia-smi，我发现每个GPU的内存限制为23040 MiB。我使用

print(f'Allocated: {torch.cuda.memory_allocated() / 1024**2} MB') print(f'Cached: {torch.cuda.memory_reserved() / 1024**2} MB')

来跟踪 GPU 的内存使用情况，发现在批量大小为 8 的情况下加载所有数据后，仅分配了大约 500 MB 的空间，并缓存了大约 2000 MB 的空间。该错误往往在调用 BERT 嵌入文本标记时发生，但稍后可能会在批量大小较小时发生。我还仔细检查了所有张量是否已正确加载到 GPU 上。

Answer 1

与 GPU 相比，您可以在 CPU 上使用更大的批量大小（但速度要慢得多），这是正常的。

一般来说，对于语言模型，每个批次使用的内存是序列长度（显然还有批次大小）的函数。如果您可以将其设置得较低，您可能不会遇到那么多内存不足问题。例如，如果您使用 Huggingface 库，您可以设置最大序列长度并使用以下方法截断：

tokenizer(text, truncation=True, max_length=128)

如果您有很长且内容丰富的文本，这当然可能会导致问题，但有时内存不足可能会在单个病态批次中出现，因为它包含比正常情况长得多的单个示例。因为每个批次（通常）都会填充到最长示例的长度，所以整个批次会变得比平常大得多。这会导致训练在整个 epoch 中途失败。