第2章 PyTorch显存管理：torch.cuda基础、显存分配与释放机制、torch.cuda.memory_summary()解读
显存管理，说白了就是跟GPU内存打交道。很多同学刚开始跑深度学习，一上来就OOM（Out of Memory），然后一脸懵。我当年也这样，第一次跑BERT-large，显存直接爆掉，还以为代码写错了。后来才发现，是我不懂PyTorch怎么管显存的。
这一章，我们就来聊聊PyTorch的显存管理。我会从最基础的torch.cuda讲起，然后深入分配与释放机制，最后教你读懂memory_summary()这个诊断利器。
2.1 torch.cuda基础：你每天都在用的API
torch.cuda是PyTorch跟GPU通信的桥梁。我个人习惯，写代码第一件事就是检查CUDA是否可用：
import torch
print(torch.cuda.is_available()) # True or False
print(torch.cuda.device_count()) # 有几张卡
print(torch.cuda.current_device()) # 当前用的是哪张卡
嗯，这里要注意：torch.cuda.is_available()返回True不代表你就能跑大模型。它只说明驱动和CUDA toolkit装好了。显存够不够，那是另一回事。
我建议你养成一个习惯——每次跑新模型前，先看看显存总量：
print(torch.cuda.get_device_properties(0).total_memory)
# 比如输出 17179869184，就是16GB
为什么强调这个？因为我在项目中遇到过，有人用RTX 3090（24GB）跑一个8GB的模型，结果还是OOM。后来发现，他把batch_size设成了128，中间变量撑爆了显存。所以，知道显存总量，才能合理分配batch size。
2.2 显存分配与释放机制：PyTorch的“小算盘”
PyTorch的显存管理，跟Python的内存管理有点像。它有一个缓存分配器（caching allocator），说白了就是：用完的显存不急着还，先留着备用。
你想想看，如果每次分配显存都去调用CUDA的cudaMalloc，那效率得多低？PyTorch的做法是：
- 分配时：先看缓存里有没有合适大小的块，有就直接用；没有才去问GPU要。
- 释放时：不真正还给GPU，而是放回缓存池，等下次用。
这个机制的好处是快，坏处是——你看着显存占用一直很高，其实很多是“假占用”。
核心概念：PyTorch的显存占用 = 实际使用 + 缓存占用。缓存占用可以通过torch.cuda.empty_cache()清空，但注意，这不会释放实际正在使用的张量。
我曾经踩过一个坑：训练完一个模型，显存占用还是90%。我以为有内存泄漏，查了半天。后来发现，只是PyTorch的缓存机制在作祟。调用一下torch.cuda.empty_cache()，显存立刻降下来了。
但要注意，empty_cache()不是万能的。它只清空缓存，不释放正在使用的张量。如果你真的想释放某个张量，得用del删除引用：
x = torch.randn(1000, 1000).cuda()
# 用完之后
del x
torch.cuda.empty_cache() # 把x占用的缓存还给系统
避坑指南：我曾经在循环里反复创建张量，忘了del，结果显存越占越多，最后OOM。记住：每个.cuda()调用，都要想好什么时候释放。
2.3 torch.cuda.memory_summary()解读：你的显存CT报告
如果你觉得显存管理像黑盒，那memory_summary()就是帮你打开这个黑盒的钥匙。我个人习惯，每次遇到OOM，第一件事就是跑这个函数：
print(torch.cuda.memory_summary())
输出会很长，但核心信息就几块。我帮你拆解一下：
| 字段 | 含义 | 我的解读 |
|---|---|---|
| Allocated memory | 当前实际使用的显存 | 这个值才是你真正关心的 |
| Active memory | 分配器正在使用的显存（含缓存） | 比Allocated大，因为包含未释放的缓存 |
| GPU reserved memory | PyTorch向GPU申请的显存总量 | 这个值通常接近你的显存上限 |
| Large blocks | 大块显存的数量和大小 | 如果大块太多，说明碎片化严重 |
举个例子，我跑一个ResNet-50时，输出可能是这样的：
| Allocated memory | 2.5 GB |
| Active memory | 3.2 GB |
| GPU reserved memory | 8.0 GB |
| Large blocks | 12 |
你看，实际只用了2.5GB，但PyTorch向GPU申请了8GB。这就是缓存机制在“囤货”。如果你发现GPU reserved memory接近显存上限，但Allocated memory还很低，那说明缓存占太多了。这时候可以试试empty_cache()。
小技巧：我习惯在训练脚本里加一个显存监控函数，每隔几个epoch打印一次memory_summary()。这样能及时发现显存泄漏或碎片化问题。
2.4 知识体系：一张图看懂显存管理
下面这张SVG图，帮你梳理本章的核心逻辑。从torch.cuda基础，到分配释放机制，再到memory_summary()诊断，一条线串起来。
2.5 实战：用memory_summary()诊断OOM
最后，我分享一个真实案例。有一次，我在训练一个Transformer模型，batch size设了32，结果OOM了。我第一反应是调小batch size，但后来想想，先看看memory_summary()再说。
输出显示：
Allocated memory: 6.2 GBGPU reserved memory: 23.8 GB（显卡是24GB）Large blocks: 48个
问题很明显：显存碎片化严重。48个大块，说明分配器把显存切得很碎，导致虽然总量够，但找不到连续的大块来放新张量。
解决方案不是调小batch size，而是：
- 在训练开始前调用
torch.cuda.empty_cache()，清空缓存。 - 检查代码里有没有频繁创建和销毁大张量的操作，尽量复用。
- 如果还不行，考虑用
torch.cuda.set_per_process_memory_fraction(0.9)限制最大使用比例。
嗯，最后我用了方案1和2，显存占用降到了18GB，batch size可以保持32不变。所以，遇到OOM别急着降batch size，先看看是不是碎片化问题。
我的习惯：每次写训练脚本，我都会在开头加一行torch.cuda.empty_cache()，确保显存状态干净。虽然不能解决所有问题，但能避免很多“灵异事件”。
好了，这一章就到这里。显存管理是深度学习优化的基本功，理解了PyTorch的分配释放机制，你就能更好地控制显存使用。下一章，我们会聊更高级的话题——梯度检查点和混合精度训练，敬请期待。