Day 42今天的主题是图像数据与显存。前几天我把一个分类模型的输入从 224×224 换到 512×512同一个脚本直接报 CUDA out of memory。当时我第一反应是“图片变大了呗”但真正排查下来发现事情远没有这么简单显存里不仅仅装着图片张量还有模型权重、梯度、优化器状态、中间激活值甚至数据加载器里一个不起眼的默认参数都能悄悄吃掉几百 MB。这篇就把今天从头到尾的排查思路、计算公式和能直接抄的优化方案整理出来给同样在“显存爆炸”边缘反复横跳的朋友一个参考。1. 图像数据在 GPU 里是怎么“排兵布阵”的1.1 一张图片从硬盘到显存要经历多少步很多人以为“把图片放 GPU 上”就是一个.cuda()的事其实在张量落到显存之前图像早就被转换了好几次。以 PyTorch 为例一张 JPEG 图片从左到右大致要经历这几个阶段读取文件PIL 或 cv2 把磁盘上的 JPEG 解压成 HWC 排列的像素数组此时一般是 uint8范围 0 到 255每个通道一个字节。数据增强与转换RandomResizedCrop、RandomHorizontalFlip等等这些操作默认在 CPU 上跑得到还是 uint8 或 float32 的图像。归一化与通道调整ToTensor()会把 HWC 变成 CHW把数值从 0 到 255 缩放到 0 到 1此时类型变成 float32单个像素占 4 字节数据量直接翻了 4 倍。送入 GPUDataLoader 把一批图像组合成 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_elementPyTorch 默认 float32bytes_per_element 4。三通道彩色图channels 3灰度图channels 1。我实际算过几个常见尺寸batch size 取 32float32结果如下输入尺寸单张像素数一个 batch 的显存占用224×224ImageNet 常见50,176约 18.4 MB512×512262,144约 96 MB1024×10241,048,576约 384 MB1920×10802,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×100200 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 提供了一组很好用的 APIimport 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(deviceNone, abbreviatedTrue))它会列出当前占用最大的一些张量、缓存块数量和碎片情况。我在一个项目里就靠这个定位到问题输入的 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_noneTrue)几个要点autocast不是把整个模型强制变成 fp16而是让 conv、linear、matmul 这类吃显存的核心算子自动用 fp16bn、softmax 等对精度敏感的层仍然保持 fp32。这样精度损失可控省下来的又确实是主要部分。反向传播必须用scaler.scale(loss).backward()不用的话小梯度会被 fp16 舍入成 0模型长时间不更新。optimizer.step()和scaler.update()的顺序不能反。先 step 让缩放后的梯度去更新权重再 update 调整缩放系数。如果发现 loss 变成 NaN 或者长时间不下降先检查损失是否在合理范围再检查是否忘了调用scaler.update()。我的实际经验同样的 ResNet-50224×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_noneTrue)这里有两个容易写错的细节。第一loss 一定要除以accum_steps。如果不除累加了 8 次之后梯度方向等于被放大了 8 倍等效学习率直接变 8 倍模型大概率发散。第二optimizer.zero_grad()不要放在每个 batch 开头否则就把前面的梯度清掉了。应该放在一个完整累积周期结束时。用set_to_noneTrue比默认写法更快因为它只是把梯度引用置空不会去把一整块显存写 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 别忽视 DataLoaderCPU 侧的隐性显存占用很多人只盯着模型忘了数据加载也会影响显存。默认的pin_memoryFalse时CPU 张量拷到 GPU 是“走一步传一步”如果设置pin_memoryTrue会先在 CPU 上分配页锁定内存拷到 GPU 的速度更快GPU 侧的临时缓冲也更少。另外num_workers不是越大越好。workers 太多时每个进程都会保存一份自己的数据副本和内存缓冲如果有collate_fn里的大图操作CPU 内存会先爆。我在一台物理机上调到num_workers8, prefetch_factor4之后发现训练开始时 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 直接定位到单次迭代的显存峰值不需要跑完整个 epochfrom torch.profiler import profile, ProfilerActivity with profile(activities[ProfilerActivity.CUDA], profile_memoryTrue) 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_bycuda_memory_usage, row_limit15))这会按算子列出每个操作的 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 报错再手忙脚乱要好得多。 
长河君 · 2026/10/16 3:43:40 · 阅读 1.2k 图像数据与显存优化：从 CUDA OOM 到 PyTorch 工程实践
Day 42今天的主题是图像数据与显存。前几天我把一个分类模型的输入从 224×224 换到 512×512同一个脚本直接报 CUDA out of memory。当时我第一反应是“图片变大了呗”但真正排查下来发现事情远没有这么简单显存里不仅仅装着图片张量还有模型权重、梯度、优化器状态、中间激活值甚至数据加载器里一个不起眼的默认参数都能悄悄吃掉几百 MB。这篇就把今天从头到尾的排查思路、计算公式和能直接抄的优化方案整理出来给同样在“显存爆炸”边缘反复横跳的朋友一个参考。1. 图像数据在 GPU 里是怎么“排兵布阵”的1.1 一张图片从硬盘到显存要经历多少步很多人以为“把图片放 GPU 上”就是一个.cuda()的事其实在张量落到显存之前图像早就被转换了好几次。以 PyTorch 为例一张 JPEG 图片从左到右大致要经历这几个阶段读取文件PIL 或 cv2 把磁盘上的 JPEG 解压成 HWC 排列的像素数组此时一般是 uint8范围 0 到 255每个通道一个字节。数据增强与转换RandomResizedCrop、RandomHorizontalFlip等等这些操作默认在 CPU 上跑得到还是 uint8 或 float32 的图像。归一化与通道调整ToTensor()会把 HWC 变成 CHW把数值从 0 到 255 缩放到 0 到 1此时类型变成 float32单个像素占 4 字节数据量直接翻了 4 倍。送入 GPUDataLoader 把一批图像组合成 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_elementPyTorch 默认 float32bytes_per_element 4。三通道彩色图channels 3灰度图channels 1。我实际算过几个常见尺寸batch size 取 32float32结果如下输入尺寸单张像素数一个 batch 的显存占用224×224ImageNet 常见50,176约 18.4 MB512×512262,144约 96 MB1024×10241,048,576约 384 MB1920×10802,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×100200 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 提供了一组很好用的 APIimport 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(deviceNone, abbreviatedTrue))它会列出当前占用最大的一些张量、缓存块数量和碎片情况。我在一个项目里就靠这个定位到问题输入的 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_noneTrue)几个要点autocast不是把整个模型强制变成 fp16而是让 conv、linear、matmul 这类吃显存的核心算子自动用 fp16bn、softmax 等对精度敏感的层仍然保持 fp32。这样精度损失可控省下来的又确实是主要部分。反向传播必须用scaler.scale(loss).backward()不用的话小梯度会被 fp16 舍入成 0模型长时间不更新。optimizer.step()和scaler.update()的顺序不能反。先 step 让缩放后的梯度去更新权重再 update 调整缩放系数。如果发现 loss 变成 NaN 或者长时间不下降先检查损失是否在合理范围再检查是否忘了调用scaler.update()。我的实际经验同样的 ResNet-50224×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_noneTrue)这里有两个容易写错的细节。第一loss 一定要除以accum_steps。如果不除累加了 8 次之后梯度方向等于被放大了 8 倍等效学习率直接变 8 倍模型大概率发散。第二optimizer.zero_grad()不要放在每个 batch 开头否则就把前面的梯度清掉了。应该放在一个完整累积周期结束时。用set_to_noneTrue比默认写法更快因为它只是把梯度引用置空不会去把一整块显存写 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 别忽视 DataLoaderCPU 侧的隐性显存占用很多人只盯着模型忘了数据加载也会影响显存。默认的pin_memoryFalse时CPU 张量拷到 GPU 是“走一步传一步”如果设置pin_memoryTrue会先在 CPU 上分配页锁定内存拷到 GPU 的速度更快GPU 侧的临时缓冲也更少。另外num_workers不是越大越好。workers 太多时每个进程都会保存一份自己的数据副本和内存缓冲如果有collate_fn里的大图操作CPU 内存会先爆。我在一台物理机上调到num_workers8, prefetch_factor4之后发现训练开始时 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 直接定位到单次迭代的显存峰值不需要跑完整个 epochfrom torch.profiler import profile, ProfilerActivity with profile(activities[ProfilerActivity.CUDA], profile_memoryTrue) 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_bycuda_memory_usage, row_limit15))这会按算子列出每个操作的 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 报错再手忙脚乱要好得多。