当前位置：首页 > article >正文

PyTorch 预分配显存是什么，为什么会有PyTorch 预分配显存

article 2025/2/7 7:52:20

💡 PyTorch 预分配显存是什么？为什么 PyTorch 要预分配显存？

📌 1. 预分配显存 = PyTorch 提前申请显存来加速计算

在 GPU 训练时，PyTorch 不会每次创建 Tensor 都去直接向显存申请内存，而是：

一次性向 GPU 申请一大块显存（预分配）
把这部分显存当作缓存，在需要创建新 Tensor 时，直接从这块缓存里分配
避免频繁的显存申请和释放，提高性能

📌 2. 为什么 PyTorch 需要预分配显存？

如果不预分配，每次创建 Tensor 都向 GPU 申请内存，会有 两大问题：

速度慢
- 直接向 GPU 申请显存的操作（cudaMalloc）非常慢，而且 GPU 需要同步，影响训练效率
- 预分配的缓存 让 PyTorch 直接从已申请的显存里拿，不用频繁 cudaMalloc
容易造成显存碎片
- 频繁分配 & 释放显存，会导致显存变得零碎，即使有空闲显存，也可能无法满足新 Tensor 的需求（类似于电脑磁盘碎片）
- 预分配大块连续显存 避免显存碎片，保证训练流畅

🚀 举个例子：如果 PyTorch 不预分配显存

假设 PyTorch 不预分配显存，每次创建 Tensor 都直接向 GPU 申请：

import torch

for _ in range(10):
    x = torch.randn(10000, 10000, device="cuda")  # 每次都向 GPU 申请新显存
    del x  # 释放显存

这样每次 x 创建时，都会向 GPU 申请 & 释放显存，导致：

频繁 cudaMalloc（显存申请）和 cudaFree（显存释放），训练速度大幅下降
容易形成显存碎片，导致 OOM（显存不足）

✅ PyTorch 预分配显存的方式

一次性向 GPU 申请较大的显存块（预分配）
在这个显存块里动态管理 Tensor：
- 创建新 Tensor：从预分配的缓存里分配显存
- 删除 Tensor：不释放显存，而是把它回收到缓存中，以便后续使用
只有显存真的不够用时，才会再次申请新显存

你可以用下面的代码观察 PyTorch 预分配的显存：

import torch

print(f"初始预分配显存: {torch.cuda.memory_reserved() / 1024**2:.2f} MB")
x = torch.randn(10000, 10000, device="cuda")
print(f"创建 Tensor 后: {torch.cuda.memory_reserved() / 1024**2:.2f} MB")

del x
print(f"删除 Tensor 后: {torch.cuda.memory_reserved() / 1024**2:.2f} MB")

torch.cuda.empty_cache()  # 释放缓存
print(f"调用 empty_cache() 后: {torch.cuda.memory_reserved() / 1024**2:.2f} MB")

可能的输出

初始预分配显存: 0.00 MB
创建 Tensor 后: 512.00 MB  # 预分配了一大块显存
删除 Tensor 后: 512.00 MB  # 仍然占用，因为 PyTorch 缓存了显存
调用 empty_cache() 后: 0.00 MB  # 彻底释放缓存

PyTorch 预分配了 512MB 显存（即使 x 只占用了 381MB）
删除 x 之后，显存仍然被 PyTorch 缓存着
调用 empty_cache() 之后，显存才真正释放

💡 总结

PyTorch 预分配显存，是为了减少频繁的显存申请 & 释放，提高训练效率
创建 Tensor 时优先从缓存分配，而不是直接向 GPU 申请
删除 Tensor 只是回收到缓存，不会立刻释放显存
empty_cache() 只是清理缓存，不会影响已存在的 Tensor

这里的显存和缓存的区别和联系是什么

显存 vs 缓存：区别和联系

🔹 显存（GPU Memory）

显存是物理存在的内存，就是 GPU 上用来存储数据和计算结果的地方。可以理解为 GPU 的“硬盘”——用来存储各种数据（比如张量、模型参数等）。

显存的作用：

存储数据：例如训练数据、模型参数、梯度、优化器状态等
用于计算：在 GPU 上进行高效并行计算的过程中，所有操作都需要使用显存
容量有限：通常 GPU 显存的容量比 CPU 内存小得多，所以需要合理管理和优化显存的使用

🔹 缓存（Cache）

缓存指的是存储在显存中的一个特殊区域，用于存放临时数据。PyTorch 的缓存机制主要是指预分配的显存，即不是每次请求显存时都向 GPU 申请，而是预先申请一块显存并用作缓存。

缓存的作用：

加速显存分配：减少频繁的显存申请（cudaMalloc）和释放（cudaFree），避免显存碎片
管理显存：PyTorch 会把一些不再使用的内存标记为“可回收”，并将其存储在缓存中，可以复用这些缓存内存来存放新的 Tensor

🔹 它们的区别：

显存是物理资源：显存本身是 GPU 上的物理内存。
缓存是显存的一部分：缓存是显存在 PyTorch 里的一部分，用来暂存已经分配的内存块，以便重复使用。

简单来说，显存就像 GPU 的内存空间，而缓存则是这块内存空间的一部分，用来加速内存管理。

🔹 它们的联系：

缓存存在于显存中，是显存的一部分。
当 PyTorch 需要显存来创建新的 Tensor 时，它会先查看缓存中是否有可用空间，如果有，就直接从缓存分配显存，避免向 GPU 申请新的显存（cudaMalloc）。
缓存会占用显存的一部分空间，而这个空间会在不再需要时被 PyTorch 回收并用于未来的 Tensor 分配。

⚡ 举个例子：

假设你正在训练一个深度学习模型，使用的是显卡（GPU）。

显存分配：
- 当你创建一个新的 Tensor，GPU 会从显存中分配内存空间。
- 如果直接从显存中申请内存，这个过程可能会很慢（尤其是在显存非常紧张时）。
缓存管理：
- 为了提高效率，PyTorch 会在 GPU 上预分配一块内存区域作为缓存。
- 当你创建 Tensor 时，PyTorch 会首先检查缓存区是否有足够的空闲空间。如果有，就直接从缓存中分配内存，而不是重新向 GPU 申请显存，这可以避免显存碎片并提高速度。
- 当你删除不再使用的 Tensor 时，PyTorch 会将这部分内存回收到缓存区，而不是立即释放给 GPU。
显存清理：
- torch.cuda.empty_cache() 会清理PyTorch 的缓存，但是它不会影响已经创建的 Tensor。
- 也就是说，缓存清理不会影响正在使用的显存，它只会释放那些不再被使用的缓存内存区域。