1. 显存优化概述：为什么需要显存优化与 Batch Size 影响

说实话，刚入门深度学习那会儿，我总觉得显存是够用的。直到有一次，我试着用一张 12GB 的卡跑 BERT-Large，训练还没跑完一个 epoch，就 `CUDA out of memory` 了。那一刻我才真正意识到——显存，就是深度学习的命根子。你想想看，模型参数、梯度、优化器状态、每一层的激活值，全都挤在显存里。模型越来越大，显存却就那么多，不优化怎么行？

为什么需要显存优化

原因其实就三个字：不够用。但细拆开来看，它背后关系到成本、效率和可能性。

  • 成本硬约束：一块 A100 80GB 的卡要好几万，多买几块成本飙升。优化显存，意味着你可以用更少的卡跑更大的模型，省下来的都是真金白银。
  • 可训练规模：同样的显存上限，优化后你能塞下更大的 batch size、更深的网络、或者更高的分辨率。我见过一个项目，就因为激活重计算，batch size 从 8 翻到了 32，收敛速度肉眼可见地快了。
  • 调试与迭代：显存溢出直接打断训练流程，重新调参、重启任务的成本很高。优化到位了，你就能安心做实验，而不是老跟 OOM 斗智斗勇。

一句话概括：显存优化不是为了炫技，而是为了让你的模型跑得起来、跑得快、跑得省。

Batch Size 对显存的影响

Batch size 是决定显存占用的头号变量。你可能觉得它只影响激活显存，其实它几乎牵动全局。我习惯从四部分来算账：

显存类别 是否受 Batch Size 影响 解释
模型参数 否 参数固定，batch size 多大多小该占多少还是多少
梯度 / 优化器状态 否 梯度也是参数尺寸，batch size 不改变
激活值（中间特征） 是 每一层输出的特征图大小与 batch size 线性正相关
临时缓冲区 部分是 像 cuDNN workspace、attention 计算中间量，与 batch size 强相关

看到没？参数和梯度占的显存基本是常量，但激活值和临时缓冲区是跟 batch size 线性增长的。举个例子，一个 ResNet-50，batch size 从 32 涨到 64，显存占用可能从 8GB 飙到 14GB —— 几乎翻倍。我曾在项目里试着把 batch size 设成 128，结果直接 OOM，后来老老实实降到 64 才跑起来。

Batch Size 对训练的影响（不止显存）

批大小的影响可不只在显存上。它还会左右梯度噪声、收敛速度和最终泛化性能。我个人经验是这样的：

  • 小 batch size（2～32）：梯度噪声大，相当于给优化加了随机扰动，有时能跳出局部极小。但训练波动大，时间也长。
  • 中等 batch size（32～256）：最常用，显存和效率平衡得不错。我一般先从这个区间开始试。
  • 大 batch size（256+）：一次训练吞吐量极高，但梯度方差小，容易收敛到 sharp minima，泛化反而可能变差。我曾经有一次把 batch size 堆到 1024，验证集准确率掉了 2 个百分点，后来加了学习率 warmup 才拉回来。

我的建议：如果你发现大 batch size 训练效果不好，先不要急着减小，试试调整学习率策略。比如用线性缩放规则：学习率 = 基础学习率 × (batch_size / 256)。这个公式在我手头好几个任务上都管用。

我曾经踩过的坑：以为 Batch Normalization 的 running mean / var 跟 batch size 无关就随便改。结果换大 batch size 后，BN 的计算反而变得更加不稳定，训练直接发散。后来我固定了 BN 的 momentum 参数才解决。

显存优化的常见思路

既然 batch size 是显存的大户，但你又不能随便砍它（不然训练效率太低），那有没有别的手段？当然有。这里简单列几个我经常用的技术：

  1. 梯度累积：前向传播时不立即更新，多积累几个 mini-batch 的梯度再一起更新。这样你能用很小的 batch size（甚至 1）跑前向，等效 batch size 却能很大。
  2. 混合精度训练（AMP）：用 FP16 保存激活和梯度，内存直接减半。我在项目中试过，训练速度提升约 40%，显存减少几乎一半，精度几乎无损。
  3. 激活重计算（Checkpointing）：前向时只保留部分激活值，反向传播需要时再重新计算一趟。这是典型的用时间换空间——多花一次计算，省下一大块显存。
  4. 模型并行 / 张量并行：把模型切到多张卡上，每张卡只存一部分参数和中间结果。适合那些单卡塞不下的超大模型。
  5. 显存碎片整理：Python 的显存分配器有时会产生碎片，尤其是在动态图里。PyTorch 的 torch.cuda.empty_cache() 虽然能释放缓存，但治标不治本。更根本的做法是尽量复用内存池，或者用 torch.cuda.set_per_process_memory_fraction 做硬限制。

一张图把这些关系梳理清楚：

显存分配与优化技术关系图 GPU 显存分配 模型参数 + 梯度 优化器状态 激活值 + 临时缓冲 Batch Size 线性影响激活值 降低显存消耗 混合精度、激活重计算 拆分/累积 梯度累积、模型并行 优化方向 核心矛盾 显存上限固定 模型规模快速增长 收益 降低成本、扩大规模 提升迭代效率 Batch Size 调节要点 小 batch → 高噪声、泛化好，但慢；大 batch → 吞吐高，但需调整学习率 先用中等 batch size（32-128）起步，观察显存和收敛曲线再调

说了这么多，其实核心就是一句话：显存优化不是让你死抠每一个字节，而是让你学会跟有限的显存共舞。把握好 batch size 这个杠杆，再搭配一两把顺手的技术（比如混合精度或梯度累积），你就能在不太改代码的情况下，把显存占用压低一大截。

嗯，这一章先到这里。我们刚才聊了为什么显存这么重要，以及 batch size 是怎么影响它的。后面会挨个深入每种优化技术的具体做法。

← 上一章 📖 返回目录 下一章 →