图像数据与显存优化：从 CUDA OOM 到 PyTorch 工程实践

Day 42，今天的主题是图像数据与显存。前几天我把一个分类模型的输入从 224×224 换到 512×512，同一个脚本直接报 CUDA out of memory。当时我第一反应是“图片变大了呗”，但真正排查下来发现事情远没有这么简单：显存里不仅仅装着图片张量，还有模型权重、梯度、优化器状态、中间激活值，甚至数据加载器里一个不起眼的默认参数，都能悄悄吃掉几百 MB。这篇就把今天从头到尾的排查思路、计算公式和能直接抄的优化方案整理出来，给同样在“显存爆炸”边缘反复横跳的朋友一个参考。

1. 图像数据在 GPU 里是怎么“排兵布阵”的

1.1 一张图片从硬盘到显存要经历多少步

很多人以为“把图片放 GPU 上”就是一个.cuda()的事，其实在张量落到显存之前，图像早就被转换了好几次。以 PyTorch 为例，一张 JPEG 图片从左到右大致要经历这几个阶段：

1. 读取文件：PIL 或 cv2 把磁盘上的 JPEG 解压成 HWC 排列的像素数组，此时一般是 uint8，范围 0 到 255，每个通道一个字节。
2. 数据增强与转换：RandomResizedCrop、RandomHorizontalFlip等等，这些操作默认在 CPU 上跑，得到还是 uint8 或 float32 的图像。
3. 归一化与通道调整：ToTensor()会把 HWC 变成 CHW，把数值从 0 到 255 缩放到 0 到 1，此时类型变成 float32，单个像素占 4 字节，数据量直接翻了 4 倍。
4. 送入 GPU：DataLoader 把一批图像组合成 batch 张量，再通过.to(device)拷进显存。

这里面最容易忽略的是第 3 步。一张 1920×1080 的普通照片，uint8 时占 1920×1080×3 ≈ 6.2 MB，但转成 float32 之后就变成 24.9 MB。训练时模型输入普遍是 float32，这个翻倍是你无法绕开的，但也正因为如此，很多“显存减肥”手段都围绕“能不能用更小的数据类型”做文章。

1.2 一张图到底吃多少字节，先学会这道公式

显存占用最核心的公式其实就一行：

bytes = batch_size × channels × height × width × bytes_per_element

PyTorch 默认 float32，bytes_per_element = 4。三通道彩色图channels = 3，灰度图channels = 1。

我实际算过几个常见尺寸，batch size 取 32，float32，结果如下：

输入尺寸	单张像素数	一个 batch 的显存占用
224×224（ImageNet 常见）	50,176	约 18.4 MB
512×512	262,144	约 96 MB
1024×1024	1,048,576	约 384 MB
1920×1080	2,073,600	约 759 MB

注意这个增长曲线：分辨率每边翻一倍，也就是从 224 到 512 再到 1024，像素数不是翻倍，而是翻 4 倍，显存占用也对应翻 4 倍。很多人在小尺寸上跑得好好的，一换大图就崩，本质就是这个二次增长规律。

这个公式的不便之处在于它只算清了“输入图像”本身。实际训练时输入张量只占显存里很小一部分，大头往往在后面。

2. 显存里不只有图片：拆解一次反向传播的完整开销

2.1 训练状态下，你的显存被谁瓜分了

我把一次常规训练迭代里显存的去向分成五块：模型参数、梯度、优化器状态、中间激活值、CUDA 初始化上下文。

模型参数好理解，所有权重加起来是多少就是多少。一个 25M 参数的模型，float32 权重大约 100 MB，模型权重本身不是最吓人的。真正容易失控的是后面几项。

梯度通常和参数同尺寸，训练时 PyTorch 会给每个需要求导的参数分配一块同样大小的梯度存储。softmax 层、全连接层的梯度还不算大，卷积层的梯度也是按参数数量来的，所以这块基本等于“再买一份模型”。

优化器状态才是隐藏大户。SGD + momentum 需要给每个参数多存一份动量，Adam 更夸张，要存一阶动量和二阶动量两份状态。也就是说，如果你用 Adam 训练一个 25M 参数的模型，光优化器状态就是 2×100=200 MB。很多 12 GB 显存的朋友训到一半爆显存，不是模型太大，而是“模型 + 梯度 + 优化器状态”这三个一模一样的拷贝把显存叠上去了。

中间激活值是最难预估的一项。正向传播时，网络每一层为了计算反向传播的梯度，都会把输入特征图保存在显存里。一个 224×224 输入、batch 32 的 ResNet 类模型，激活值轻松超过 1 GB；如果把分辨率提到 512，激活值可能直接冲向 4 GB 以上。这也是为什么工程上常说“激活值决定了显存天花板”。

最后还有一块容易被忽略的 CUDA context。只要你在程序里第一次调用 GPU，驱动就会预留几百 MB 空间给 kernel 调度、上下文管理器等，这和你模型多大没关系。所以显存明明 8 GB，实际你能支配的往往只有 7.5 GB 左右。

2.2 一张图告诉你 fp32、fp16、bf16 的差距

解决显存问题，大多数优化手段最后都落在“数据类型”上。float32 占 4 字节，float16 和 bfloat16 占 2 字节。单看一个数字，同样的 batch 和分辨率，半精度输入张量的显存直接腰斩，激活值也跟着腰斩，模型参数和梯度也同样减半，只有优化器状态里的 fp32 主权重还是 4 字节。

这就是为什么混合精度训练几乎是现代训练的标配，它不是为了“变快”这么简单，更是为了“省显存”。后面我会单独讲实操配置。

不过这里注意一个反直觉的点：显存里经常同时存在 fp16 和 fp32 两份权重。AMP 训练时，模型前向用 fp16 计算，但优化器里必须保留一份 fp32 的“主权重”，否则小梯度会被半精度舍入掉。所以混合精度省的是激活值和中间计算量，模型权重那块并不会像想象中那样缩水一半。

3. 实操记录：我如何把 batch size 从 16 提到 128 还不爆显存

3.1 动手优化之前，先给项目做一次显存“体检”

我踩过的坑之一就是不看数据就调参数。后来养成了习惯：任何模型跑起来之前，先记录训练过程中的显存峰值。PyTorch 提供了一组很好用的 API：

import torch torch.cuda.reset_peak_memory_stats() # 正常训练循环... memory_used = torch.cuda.max_memory_allocated() / 1024**3 print(f"显存峰值: {memory_used:.2f} GB")

reset_peak_memory_stats()会把峰值计数器归零，max_memory_allocated()返回从归零时刻到现在的最大分配量。它比 nvidia-smi 更精确，因为 nvidia-smi 看到的是整个进程的占用，包括缓存碎片，而 PyTorch 的统计能直接告诉你实际张量占了多少。

只看数值还不够，建议配合 PyTorch 官方 memory summary 看细化分布：

print(torch.cuda.memory_summary(device=None, abbreviated=True))

它会列出当前占用最大的一些张量、缓存块数量和碎片情况。我在一个项目里就靠这个定位到问题：输入的 batch 本身只占 300 MB，但某个中间特征图因为尺寸计算错误被放大到了 8,000 张，瞬间把显存打满。这种问题光看 nvidia-smi 完全看不出来，必须看内存摘要。

3.2 方案一：混合精度训练，一行配置省一半激活值

如果是 PyTorch 1.6 以后的版本，直接使用官方自带模块，不需要额外装第三方库：

from torch.cuda.amp import autocast, GradScaler scaler = GradScaler() for images, labels in dataloader: images = images.cuda() # 自动在算子层面选择 fp16 或 fp32 with autocast(): outputs = model(images) loss = criterion(outputs, labels) # 用 scaler 反向传播，防止梯度下溢 scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() optimizer.zero_grad(set_to_none=True)

几个要点：

• autocast不是把整个模型强制变成 fp16，而是让 conv、linear、matmul 这类吃显存的核心算子自动用 fp16，bn、softmax 等对精度敏感的层仍然保持 fp32。这样精度损失可控，省下来的又确实是主要部分。
• 反向传播必须用scaler.scale(loss).backward()，不用的话小梯度会被 fp16 舍入成 0，模型长时间不更新。
• optimizer.step()和scaler.update()的顺序不能反。先 step 让缩放后的梯度去更新权重，再 update 调整缩放系数。
• 如果发现 loss 变成 NaN 或者长时间不下降，先检查损失是否在合理范围，再检查是否忘了调用scaler.update()。

我的实际经验：同样的 ResNet-50，224×224、batch 64，开启 AMP 后显存峰值从原来的 4.2 GB 降到 3.1 GB，大概省了 25% 到 30%。别期待腰斩，因为优化器状态和主权重还在。

3.3 方案二：梯度累积，用“时间”换“空间”

如果你想维持 batch size 128，但显存只装得下 batch 16，梯度累积是最直接的方案。做法是：每 8 个小 batch 才更新一次参数，等效的“经验上”batch size 还是 128。

accum_steps = 8 # 128 / 16 optimizer.zero_grad() for step, (images, labels) in enumerate(dataloader): images = images.cuda() with autocast(): outputs = model(images) loss = criterion(outputs, labels) / accum_steps # 关键：先除以累加步数 scaler.scale(loss).backward() # 梯度会累加到已有的 gradient 里 if (step + 1) % accum_steps == 0: scaler.step(optimizer) scaler.update() optimizer.zero_grad(set_to_none=True)

这里有两个容易写错的细节。

第一，loss 一定要除以accum_steps。如果不除，累加了 8 次之后梯度方向等于被放大了 8 倍，等效学习率直接变 8 倍，模型大概率发散。

第二，optimizer.zero_grad()不要放在每个 batch 开头，否则就把前面的梯度清掉了。应该放在一个完整累积周期结束时。用set_to_none=True比默认写法更快，因为它只是把梯度引用置空，不会去把一整块显存写 0。

梯度累积的原理就是用更长的迭代换更低的瞬时显存。我一般建议把算力利用率放第二位，先保证模型能跑起来。等正常训练后，再逐步提高 batch size，把累积步数降下来。

3.4 方案三：激活值重计算，把 3 GB 中间值压到 300 MB

如果输入尺寸不能变，batch size 又想维持，激活值就是最大的瓶颈。这时候可以用激活重计算，也叫 gradient checkpointing。

PyTorch 里已经有封装好的接口：

from torch.utils.checkpoint import checkpoint # 自定义网络时，把关键层包进 checkpoint def forward(self, x): x = checkpoint(self.encoder_block, x) x = self.decoder_block(x) return x

原理是：正向传播时不再保存每一层的激活值，只保留最后一层结果和输入；反向传播求梯度时再临时重算一遍这些激活值。空间省了，多付出的是大约 1.3 到 1.5 倍的前向计算时间。

实际项目中，我不建议把每个模块都包进去。优先包那些输出特征图特别大的层，比如下采样之前的大分辨率阶段。小模块包进去收益不大，还平添开销。

3.5 别忽视 DataLoader：CPU 侧的隐性显存占用

很多人只盯着模型，忘了数据加载也会影响显存。默认的pin_memory=False时，CPU 张量拷到 GPU 是“走一步传一步”；如果设置pin_memory=True，会先在 CPU 上分配页锁定内存，拷到 GPU 的速度更快，GPU 侧的临时缓冲也更少。

另外，num_workers不是越大越好。workers 太多时，每个进程都会保存一份自己的数据副本和内存缓冲，如果有collate_fn里的大图操作，CPU 内存会先爆。我在一台物理机上调到num_workers=8, prefetch_factor=4之后，发现训练开始时 GPU 显存被多个预取缓冲挤爆。这不是模型问题，是 DataLoader 预取数据超过了 GPU 消化速度。解决办法是把prefetch_factor调小到 2，或者让 workers 减少。

4. 常见问题与排查技巧实录

4.1 训练一半突然 OOM：先分清是“模型问题”还是“进程残留”

训练到第 5 个 epoch 突然报 CUDA out of memory，很多人第一反应是调小 batch size。但我的排查习惯是，先看是不是上一个 Python 进程还没死。

“进程残留”是显存排查里最常见也最闹心的问题。跑完一个大实验后直接关掉终端，进程可能没被 kill 干净，显存还被占着。新的进程一启动，分到的可用显存只剩很小一块。排查命令很简单：

nvidia-smi

看 GPU 列表里的进程 PID。如果有不属于当前程序的进程，直接用kill -9 PID杀掉。

还有一种情况：同一个进程里连续加载了好几个模型，之前的模型变量还挂在某个 list 或 dict 里没释放。这种我能用的排查办法是检查代码里有没有把net、model之类的大对象保存为全局变量，或者在循环里反复实例化模型。

4.2 为什么我调用了 empty_cache() 显存还是没降

torch.cuda.empty_cache()经常被误解成“清空显存”。它实际上只是把 PyTorch 缓存分配器里尚未归还给驱动的那部分可释放块返回给 CUDA。如果你在nvidia-smi里看到的进程占用没有变化，说明显存里大多数块是正在使用的 tensor，根本不在可释放范围内。

更常见的现象是：一个训练循环里显存持续增长，每步都在涨，涨到最后 OOM。这大概率不是模型参数在涨，而是某个 tensor 被反复创建但引用没断，或者是模型在训练和验证模式下反复切换时，激活值计算图没有正确释放。

我的处理方式是：验证阶段一定加上with torch.no_grad()，并且调用model.eval()。少了no_grad()，验证时 PyTorch 也会构建计算图、保存激活值，显存可能直接翻倍。这是一行代码的差别，却是我见过最多人踩的坑。

4.3 从 224 换到 512 就崩：二次方增长不是开玩笑

开头说的那个案例，我从 224 切到 512 就 OOM，直接用公式算一下就明白了。假设当时 batch 是 64：

• 输入张量：64×3×512×512×4 = 201 MB
• 第一层卷积后的激活图已经变成 64×64×256×256 这个量级
• 如果网络有 4 个 stage，激活值基本是线性往上加，轻松超过 2 GB
• 再加上模型参数、梯度、Adam 状态，总共 4 GB 到 6 GB 是很正常的

所以如果你是从小图迁移到大图，我建议按“先降 batch，再降分辨率，最后降模型宽度”的顺序调整。先在 512 下用 batch 8 跑通，确认显存有余量，再逐步往上加 batch。不要一次性把 batch 保持原样往上怼。

如果你确实需要高分辨率输入，比如做检测或分割，更高分特征图的激活值更夸张。此时前面说的 checkpoint 几乎是必须开的。还有一个小技巧：让图像在输入网络之前先插值到一个合理尺寸，很多增强库会把超大数据直接送进网络，这一步就能省掉几个 GB。

4.4 我每次调 batch 都要重启训练，有没有更快的方法继续

调试显存时反复重启训练确实浪费时间。我后来会用 PyTorch 的 memory profiler 直接定位到单次迭代的显存峰值，不需要跑完整个 epoch：

from torch.profiler import profile, ProfilerActivity with profile(activities=[ProfilerActivity.CUDA], profile_memory=True) as prof: for images, labels in dataloader: images = images.cuda() with autocast(): outputs = model(images) loss = criterion(outputs, labels) scaler.scale(loss).backward() prof.step() break print(prof.key_averages().table(sort_by="cuda_memory_usage", row_limit=15))

这会按算子列出每个操作的 CUDA 显存使用量。我至少两次靠它找到“瓶颈算子”，而不是靠猜。最典型的是某个自定义前处理层把中间 tensor 放大，或者是把整张图在 GPU 上做了多次 resize。

4.5 几个反直觉但很实用的显存小习惯

• 输入张量尽量在 CPU 上完成归一化和 resize，不要在 GPU 上做。ToTensor和Normalize放在 DataLoader 的 transform 里，比放在 forward 前更省显存，因为 GPU 上少了一块临时 float 中间量。
• 如果你用的是 Transformers 类模型，序列很长时激活值按序列长度的平方增长，这时候 padding 策略要特别小心。把训练样本按长度分桶，或者用torch.utils.data.BatchSampler做相似长度分组，能显著降低同 batch 内的 padding 浪费。
• 在训练脚本开头加上torch.backends.cudnn.benchmark = True，虽然它不直接省显存，但能加速卷积搜索，让同样 batch 下的迭代更快，等于变相提高了显存利用率。
• 集群环境下，显存不够不一定要换 GPU，可以用torch.utils.checkpoint搭配梯度累积，把单卡能做满的实验做满，再考虑多卡。

5. 这 42 天下来，我的一点体会

把 Day 1 到今天的问题串起来，我发现“图像数据与显存”其实是深度学习中第一个真正需要“算账”的地方。你没法凭感觉决定 batch size，也没法凭感觉决定输入分辨率。每次调参之前，先用公式算一遍输入张量有多大，再用 memory summary 确认激活值占了多少，心里有数之后再去动网络结构或优化器配置，方向基本不会错。

我个人现在养成的固定流程是：先写一个最小复现脚本，打印max_memory_allocated()，确认单 batch 的显存占用；然后开 AMP，用梯度累积把有效 batch 撑到目标值；如果还差一点，再对重模块开 checkpoint。这个顺序能解决我遇到过的 90% 以上显存问题。

最后再分享一个小技巧：把显存峰值记录写进训练日志里，每个 epoch 末打印一次。不要只看第一次，随着训练进行，某些层的激活模式会因输入分布变化而波动，显存占用也会起伏。盯着这条曲线，比等到 OOM 报错再手忙脚乱要好得多。