最近在跑一个简单的ResNet50迁移学习项目,数据集大概2万张图片,每张resize到224x224。我设batch size=32,结果跑第一个epoch就报CUDA out of memory。我用的是RTX 3060 12G,按理说应该够用吧?是不是我DataLoader里开了太多num_workers?还是说模型里用了什么隐藏的显存泄露?尝试把batch size降到8倒是能跑了,但训练慢得离谱,而且准确率也不太行。
想请教下各位大佬,这种场景下一般怎么优化显存?混合精度、梯度累积这些方法真的能立竿见影吗?还是说我哪里写得不规范?求指点,孩子快被显存整自闭了。
用PyTorch做图像分类训练时显存炸了,是我代码写错还是batch size太大?
全部回复
共 152 条30系卡对cudnn和pytorch版本挺敏感的,我之前也是3060 12G跑ResNet50 batch size设24都炸,降到16就好,你这32确实有点莽。混合精度大概率能救,半精度显存直接砍半,配合梯度累积效果更稳,不过注意bn层要调成float32。另外检查下dataloader的pin_memory是不是开了,某些情况下它也会偷偷吃显存。准确率低可能是lr没跟着batch size调,试试按比例缩放一下。
12G显存跑resnet50 batch size=32按理说是够的,我自己的3060跑过类似的,batch size开到64才爆。你这情况大概率不是batch size的问题,先检查一下输入图片的尺寸——你说resize到224x224,但dataloader里是不是忘了加transform?如果原始图片没被正确resize,直接拿原图尺寸喂进去,显存瞬间就炸了。我之前踩过这个坑,加载的图片是3000x4000的,resize写在预处理里但没应用到dataset上,跑第一个batch直接崩。
另外num_workers开太多确实会吃显存,但不是主要原因,因为worker是在CPU上加载数据,除非你用了pin_memory=True并且进程数开到了8以上,才会把部分数据锁在显存里。你试试把num_workers降到2或者4,pin_memory先关掉看看。
混合精度和梯度累积确实立竿见影。用torch.cuda.amp.autocast加上GradScaler,基本能省30%-40%显存,而且训练速度还快。梯度累积的话,你batch size=8,梯度累积步数设4,效果相当于batch size=32,显存占用只相当于8。不过要注意学习率要相应调大,否则收敛慢。
还有个小技巧——检查一下模型里是不是用了dropout或者batchnorm的training模式没切对,有时候eval模式下某些层的缓存也会吃显存。跑第一个epoch就崩,不如先试试单张图片能不能跑通,逐步加batch size定位到临界点,别一上来就莽。
我也是刚入坑PyTorch不久,看到你这情况简直世另我。我用的也是3060 12G,跑ResNet50 batch size=32按理说真不算大,之前我跑个更轻量的EfficientNet-B0都差点爆,后来发现问题是出在DataLoader的pin_memory=True上,关掉之后显存占用直接降了快2G。你可以试试把这个选项关了,或者pin_memory=True但num_workers别开太高,我一般设4就够用了,太高反而容易让显存碎片化。
另外你说的混合精度和梯度累积,我亲自试过,确实有效。混合精度用torch.cuda.amp,基本无痛省一半显存,而且训练速度还快一点,就是得注意loss scaling的细节,别让梯度下溢。梯度累积的话,你设batch size=8,累积4步效果相当于batch size=32,但显存只占8张图的空间。不过要注意,如果用了BatchNorm,梯度累积可能会让BN的统计量不太准,得小心点。
还有个小坑,你是不是用了pretrained=True加载模型?有时候预训练权重里带的一些缓存或者中间变量没释放干净,可以试试先加载到CPU再转cuda。另外检查下你的dataloader里有没有在每次迭代后手动回收变量,比如用del或者清空cache,有时候一个torch.cuda.empty_cache()就能救急。
至于准确率下降,batch size变小确实会让梯度噪声变大,但如果你用梯度累积,等效batch size不变,理论上精度应该差不多。你降到8后准确率不行,会不会是学习率没跟着调?小batch size一般要稍微降点lr,或者用warmup。建议先试试混合精度+梯度累积,同时把num_workers降到4,pin_memory关掉,看看能不能稳定跑起来。要是还不行,可以贴一下代码片段,大家帮你看看是不是有显存泄露。
12G显存放不下batch size=32的ResNet50+224x224输入,这其实挺正常的,不用太怀疑自己。你看到的那些“3060 12G随便跑”的说法,多半是别人用的输入尺寸更小或者模型没加载完整,或者干脆没把验证集的显存算进去。ResNet50单张224x224的前向大概占200-300MB,但反向传播要翻倍,加上Adam优化器的动量缓存、DataLoader的预加载、以及PyTorch的cuda缓存分配策略,实际峰值很容易冲到10G以上,32张一压就炸。
你降到8能跑但准确率不行,这个锅batch size不背。准确率下降大概率是学习率没跟着调,或者BN层在batch size=8时统计量抖动太大。建议你保持batch size=8,把学习率从默认的0.1/0.01按比例降到0.025左右,同时把BN换成GroupNorm或者用SyncBN(单卡也能用,但效果有限)。梯度累积确实有用,但本质只是把显存压力分散到时间上,你设accumulation_steps=4,等效batch size=32,但每步更新时显存峰值还是8张的量,不会额外炸。
混合精度是立竿见影的,AMP开了之后显存直接打六折,12G跑32张绰绰有余。但注意amp的GradScaler要配合梯度裁剪,不然容易loss NaN。另外检查下你的DataLoader里pin_memory=True和num_workers=4-8是标配,不会导致显存泄露,除非你手动在worker里创建了tensor并保持引用。还有一个隐蔽的点:如果你在训练循环里每步都新建了tensor又没有及时del,或者用了requires_grad=True的变量做in-place操作,可能会触发显存碎片化,表现为逐步上涨直到崩。你可以用torch.cuda.memory_summary()看一下是分配失败还是碎片太多。
最后建议你先把batch size改回32,开amp,关掉任何自定义的hook或日志记录,跑一个epoch试试。如果还炸,大概率是模型里某些层(比如自定义的注意力模块)没有释放中间变量。3060 12G跑迁移学习ResNet50绝对够用,别自闭。
3060 12G跑ResNet50 batch size 32按理说确实不该炸,我拿同款卡试过,ResNet50输入224x224,fp32下batch 32大概占用9.5G左右,还有余量。但你2万张图跑第一个epoch就炸,大概率不是batch size的问题,而是DataLoader那边搞了什么骚操作。
先检查几个地方:第一,你在DataLoader里是不是开了pin_memory=True,然后num_workers又设得很高?比如设成8或者12?这玩意在windows上特别容易爆显存,因为worker进程会复制一部分数据到显存里。建议pin_memory先关掉,num_workers先降到2或者4试试。第二,如果你用了albumentations或者torchvision的随机增强,检查一下是不是在__getitem__里把原图和增强图都保留下来了,没及时释放。第三,模型定义里有没有不小心用了.eval()或者train()之外的奇怪hook,比如把中间层的feature map全存下来做可视化?
混合精度(AMP)对显存优化很直接,fp16能把显存占用砍半还多,而且现在torch.cuda.amp基本无感接入,几行代码的事,推荐先上这个。梯度累积是治标不治本,它只是让你用更小的batch模拟大batch的梯度,但显存占用还是小batch的水平,训练慢是肯定的。你batch降到8之后准确率不行,可能是因为batch norm的统计量在小batch下不稳定,可以试试把syncBN打开,或者干脆用大batch+梯度累积来模拟。
另外还有个骚操作:把图片尺寸先缩到160x160做预训练,最后几个epoch再切回224,显存能省一大截,而且最终精度几乎不受影响。我之前在类似项目里这么干过,训练时间直接减半。
最后说一句,显存泄露在PyTorch里其实挺常见的,尤其是用了自定义Dataset还手动调用了.cuda()或者.to(),建议检查一下验证集或者测试集的数据加载逻辑,是不是每次迭代都往显存里塞新张量但没有清掉旧的。用torch.cuda.empty_cache()清理一下缓存,或者干脆用with torch.no_grad()包住推理部分。
3060 12G跑224的ResNet50,batch size=32确实有点紧,但也不至于直接炸,排查下是不是开了pin_memory=True加上num_workers太高,IO线程把显存占了一部分。混合精度(autocast+gradscaler)对这类场景改善非常大,基本能省40%显存,训练速度还更快,梯度累积也能救急但会拖慢收敛。另外检查下有没有不小心把验证集的梯度也打开了,很多人会在val loop里漏掉torch.no_grad()。
12G显存跑224分辨率、batch size 32的ResNet50按理说不会直接炸,你这大概率不是batch size的问题。检查下DataLoader里的pin_memory是不是开了,以及transform里有没有无意中把图片转成FP32再喂进网络。混合精度(AMP)对这类场景基本是立竿见影的,显存能省将近一半,配合梯度累积还能等效增大batch size。另外建议先用torch.cuda.memory_summary()看下哪块占得最多,有时候是优化器状态或中间变量没释放。
3060 12G跑ResNet50 batch size=32按理说确实不该炸,我3070 8G跑32有时候都能撑住。你这个问题大概率不是batch size本身的锅,建议先排查几个点。
第一,检查一下DataLoader里图片加载完是不是转成float16或者uint8了?很多人习惯用torchvision的transforms,但忘了ToTensor()会把PIL图像转成float32,如果前面没做归一化或者类型转换,显存会多占一倍。第二,num_workers开太多确实会占显存,每个worker会预加载一批数据到显存里,建议先设成0或者2试一下,排除这个因素。第三,模型本身有没有开eval模式或者用了dropout?ResNet50预训练模型如果没冻结BN层,训练时显存也会额外消耗。
混合精度(AMP)在3060上效果很明显,我跑类似任务显存直接砍半,而且速度还快了。梯度累积本质上是时间换空间,batch size降到8的话,梯度累积4步等效batch size=32,但显存占用还是8,只是能模拟大batch的效果,对准确率没坏处。不过你提到准确率不行,可能是学习率没同步调整,batch size小了lr也该相应降一降。
另外建议看看nvidia-smi里显存有没有被其他进程占着,有时候pytorch的缓存没清干净也会这样。可以先跑个最小的测试脚本,只加载模型和一张图,看看基线显存是多少,再逐步加batch和worker,定位到具体哪一步炸的。别自闭,这种问题调几次就有经验了。
这个问题其实挺典型的,属于那种“看着配置够,跑起来就炸”的日常。RTX 3060 12G在显存容量上确实不算小,但ResNet50 + 224x224 + batch size 32这个组合,理论上应该刚好卡在显存边缘,如果代码里再有一些隐性开销,炸掉并不意外。我先帮你拆一下可能的原因,再给一套从浅到深的优化方案。
首先,直接回答你最关心的问题:是不是batch size太大?是,也不完全是。ResNet50本身参数量约25M,FP32推理时单张224x224图片的前向显存占用大概在200-300MB左右,但训练时因为要保存中间激活值用于反向传播,这个数字会膨胀到1-2GB每张图(具体取决于实现和输入尺寸)。32张图就是32-64GB的激活显存需求,这还没算模型参数、优化器状态、DataLoader缓冲。12G显存显然扛不住。所以batch size 32对于3060 12G跑ResNet50全精度训练,确实偏大。但你说降到8能跑但准确率不行,这个“准确率不行”大概率不是batch size直接导致的,而是你学习率没跟着调,或者迁移学习时冻结层/数据增强策略没配合好。batch size从32降到8,如果学习率还是原来那套,梯度噪声变大,收敛自然会变差。后面再细说。
至于num_workers,它主要影响CPU内存和DataLoader的IO效率,和GPU显存没有直接关系,除非你用了pin_memory=True并且num_workers开得特别多(比如超过8个),可能会导致CPU侧预分配大量页面锁定内存,间接给GPU显存传输造成压力,但通常不会直接导致OOM。你可以先检查一下代码里有没有无意中把整个数据集加载到GPU上,或者自定义Dataset的__getitem__里有没有保留大张量没释放。更常见的隐性问题是:你在训练循环里每步都创建了新的tensor却没有显式删除,或者用了torch.no_grad()的地方没加对,导致计算图累积。我见过最离谱的案例是有人在每个batch里把loss.item()赋值给一个全局列表,然后列表越积越大,最后把CPU内存撑爆了,连带GPU显存也受影响。
现在说优化方案。我按见效速度从快到慢排列,你可以根据自己情况组合使用。
第一梯队:混合精度训练。这是目前性价比最高的显存压缩方法,几乎零代码改动。用PyTorch自带的torch.cuda.amp,只需要在训练循环里加一个GradScaler和autocast上下文。原理是:前向和反向时,大部分计算用FP16,但权重更新用FP32副本,这样激活显存几乎减半,同时模型精度损失可以忽略(尤其是ResNet这种对精度不敏感的大模型)。你3060的Tensor Core虽然不如30系高端卡多,但跑FP16加速和显存减半的效果依然明显。我实测ResNet50 + batch size 32在3060上,开amp后显存占用从约11GB降到6-7GB,完全能跑。你降到batch size 8时开amp,甚至可能直接上64。不过要注意:amp对某些自定义算子或者LayerNorm层可能有数值溢出风险,但ResNet50全是Conv和BatchNorm,很安全。
第二梯队:梯度累积。这个跟混合精度搭配使用更香。原理是你每个step只计算一个小batch的loss和梯度,但先不更新参数,累积几个step的梯度后再统一更新。这样等效于用更大的batch size训练,但显存占用只相当于一个小batch。比如你batch size 8跑不动,可以设batch size 2,然后累积4步,效果等同于batch size 8。但要注意:梯度累积会降低训练速度(因为每个step都要做前向+反向),而且BN层的统计量会受影响,因为每个小batch的样本数太少。解决方案是:要么用SyncBN(但需要多卡),要么把累积步数控制在4-8之间,同时把BN的momentum调大一点。对于你的2万张图片,batch size 2累积4步,等于每个有效batch看到8张图,BN统计量还算稳定。
第三梯队:检查DataLoader和数据集加载。很多人会忽略一个点:如果图片是jpg格式,torchvision的默认解码用的是PIL,会一次性把整张图片解码成RGB张量再resize。2万张图每张224x224,如果Dataset在__init__里就把所有图片读取到内存里,CPU内存会爆炸,进而影响GPU的page-locked memory分配。正确做法是:Dataset只存文件路径,在__getitem__里按需读取。如果你已经这么做了,那可以试一下ImageFolder配合transforms,并且把num_workers设为4或者8(不要超过CPU核心数的一半),pin_memory=True。另外,检查一下transforms里有没有用RandomResizedCrop这种需要多次解码的算子,如果有,可以换成固定resize + 随机裁剪,减少计算量。
第四梯队:模型层面的优化。ResNet50的参数量其实有很多冗余,你既然做迁移学习,可以考虑冻结前几层。通常冻结前3个stage(即从conv1到layer2),只训练layer3、layer4和全连接层。这样反向传播时不需要保存冻结层的激活值,显存占用能再降20-30%。具体做法:for param in model.parameters(): param.requires_grad = False,然后按需打开后几层的梯度。注意冻结层后要调低学习率,因为可训练参数变少了,学习率太大容易震荡。
第五梯队:梯度检查点(activation checkpointing)。这是终极手段,但会牺牲约15-20%的训练速度。原理是在前向时只保存部分层的激活值,反向时重新计算被丢弃的激活值。PyTorch官方有torch.utils.checkpoint接口。对ResNet50,你可以把每个BasicBlock或Bottleneck包裹在checkpoint里,这样显存占用能从每张图2GB降到0.5GB以下。代价是每个step要多做一次前向计算(反向时重算),训练时间会延长。但如果你batch size实在调不上来,这是最后的救命稻草。
第六梯队:分布式训练或混合并行。这个对你单卡场景可能不太适用,但如果你以后换多卡,可以用Zero Redundancy Optimizer (ZeRO) 或者模型并行。不过3060不支持NVLink,多卡通信效率低,不推荐。
再说你提到的准确率问题。batch size 8训练慢且准确率差,很可能是因为学习率没调。迁移学习通用的经验是:先用一个很小的学习率(比如1e-4)做几个epoch的warmup,然后逐渐增大。如果你从8张图突然跳到32张图,学习率应该相应增大(比如增大到原来的2-4倍),因为更大的batch size对应更稳定的梯度方向,需要更大的步长。反之,从32降到8,学习率应该降低。你可以尝试用cosine annealing scheduler,初始学习率设为1e-3或5e-4,配合warmup 5个epoch。另外,2万张图片对于ResNet50来说其实偏少,建议做更强的数据增强:RandomHorizontalFlip、ColorJitter、RandomRotation、CutMix或MixUp。这些增强会增加训练难度,但能提升泛化能力,同时因为样本多样性增加,对小batch size的敏感性也会降低。
最后,说一个你可能没注意到的细节:显存碎片化。PyTorch在频繁分配和释放张量时,显存会产生碎片,导致虽然总空闲显存还够,但无法分配一块连续的大内存。这种情况在训练初期尤其常见,因为DataLoader的prefetch和模型初始化会分配很多临时张量。解决办法是:在训练开始前先用一个dummy batch跑一次前向+反向,预热CUDA缓存,然后清空缓存(torch.cuda.empty_cache()),再开始正式训练。另外,设置torch.backends.cudnn.benchmark=True可以让cudnn自动选择最优卷积算法,减少临时显存分配。
总结一下你现在可以立刻做的操作: 1. 开启混合精度(amp),batch size设回32,如果还炸就降到16或24,配合梯度累积。 2. 检查Dataset实现,确保没有预加载全部图片。 3. 冻结ResNet50的前几个stage,只训练最后两层。 4. 把学习率调到1e-4左右,用cosine scheduler。 5. 如果还不行,用梯度检查点,但这是最后选项。
另外,不要迷信“12G显存够用”这个说法。实际上,RTX 3060的12G是GDDR6,带宽和频率都不如RTX 3080的10G,所以即使显存容量大,但实际吞吐量有限。你跑ResNet50,batch size 32开amp能稳定在30-40毫秒每step就算不错了。如果追求速度,可以考虑用更轻量的模型,比如ResNet34或者EfficientNet-B0,准确率在迁移学习场景下未必比ResNet50差多少,但显存占用直接减半。
最后,显存炸了不是世界末日,它只是逼着你把代码写得更好。很多人觉得显存不够就换卡,但实际上通过合理的优化,很多模型都能在有限的硬件上跑起来。你在论坛上看到那些动不动就上A100的大佬,其实他们的代码里大概率也藏着类似的优化技巧,只是不说而已。
12G跑224的ResNet50,batch32确实有点极限,试试梯度累积两步模拟32的效果。
12G跑ResNet50加224分辨率,batch size=32按理说不至于直接崩,你先检查下pin_memory是不是设成True了,这个在3060上容易爆显存。混合精度真的挺管用的,开torch.cuda.amp能省将近一半显存,配合梯度累积你甚至可以回到batch size=32的效果。另外确认下你用的不是ResNet50的官方预训练权重吧?那个带BN层有时会额外吃显存,可以试试冻结前几层。准确率不行大概率是lr没调对,batch size变小后要相应降低学习率。
12G跑32的batch size确实有点勉强,试试梯度累积2步,效果跟32差不多但显存压力小很多。
12G跑ResNet50加224尺寸,batch size=32确实有点极限,不是你的代码问题。混合精度和梯度累积立竿见影,前者直接省一半显存,后者让你用更小batch模拟大batch效果,建议优先试试这两招。另外num_workers设成4或者8就行,开太多反而容易爆显存,准确率低可能是学习率没调好,batch缩了之后记得按比例调一下lr。
12G跑ResNet50 batch32确实有点悬,224x224的图一张就占不少,加上梯度缓存和中间变量很容易爆。混合精度效果很明显,开AMP基本能省一半显存,另外把num_workers设成4或者8也够了,开太多反而会增加内存开销。你降到batch8训练慢可以试试梯度累积,模拟大batch效果,准确率低可能是学习率没跟着调,先看看这些能不能救一下。
12G跑224的ResNet50加batch32确实有点极限,尤其是迁移学习如果还带着BN层的话。我建议你试试梯度累积,相当于虚拟增大batch size但显存占用不变,对收敛也有帮助。混合精度也值得开,AMP在30系卡上收益很明显,基本不影响精度。另外检查下DataLoader的pin_memory和num_workers,设太高反而容易爆显存,4到6个workers就差不多了。
12G跑224的ResNet50 batch size到32确实有点极限,尤其如果你的图片加载没做缓存或者用了随机增强,显存波动会很大。建议先试试梯度累积,用batch size 8累积4步等效32,显存占用和训练速度都能平衡。混合精度也挺管用的,AMP开起来显存能省将近一半,而且对准确率几乎没影响。另外num_workers建议设成4或者8就行,太高了反而会占用额外的显存缓存,不是越多越好。
12G跑ResNet50加224的图,batch size=32按理说不是完全没希望,但确实有点极限,尤其是如果你开了输入输出通道的额外占用。你可以先用torch.cuda.max_memory_allocated看一眼峰值显存到底花在哪,很多情况是DataLoader的prefetch_factor或者num_workers开太高,数据预处理阶段疯狂占显存。混合精度(AMP)在这种场景下效果挺明显的,光是把float32切到float16就能省将近一半,而且对准确率基本没影响,建议你试试torch.cuda.amp。梯度累积也是个好思路,相当于用时间换空间,batch size=8累积4步就等效于32了,训练速度虽然慢点但起码能跑。另外检查下模型里是不是忘了加model.train()或者有什么不必要的梯度缓存没清,有时候显存炸是代码里偷偷保持了中间变量。至于你说准确率不行,batch size降到8后模型收敛特性本来就会变,学习率可能需要跟着调,别直接套原来的参数。最后如果实在不行,可以试试用torch.utils.checkpoint给ResNet的block做梯度检查点,牺牲一点计算换显存,3060的12G其实挺能扛的。
12G显存跑ResNet50加2万张224的图,batch size=32按理说不会直接炸,我怀疑问题出在DataLoader的num_workers设太高了,比如开到8甚至更多,这会导致每个worker预加载大量图像到显存里。你可以试试把num_workers降到2或者0,看看有没有改善,有时候这比降batch size还管用。混合精度训练确实是立竿见影的,用torch.cuda.amp把forward和loss计算包在autocast里,显存能省30%到50%,而且速度还快,我自己的3060上跑ResNet50,batch size从16直接提到了40都没问题。梯度累积也是个好办法,如果你不想降batch size,可以设accumulation_steps=4,相当于模拟batch size=32但每次只算8张的梯度,显存压力小很多。不过准确率不行可能跟batch size本身关系不大,更像是学习率没调或者数据增强不够强,毕竟迁移学习对小batch size没那么敏感。你可以先试试混合精度加num_workers=2,如果还炸,再考虑梯度累积,别急着怀疑代码写错,这种显存问题多半是配置没优化好。
12G跑ResNet50加224的图,batch size=32按理说不会直接炸,我猜大概率不是batch size本身的问题。你检查一下DataLoader里pin_memory是不是设成了True,这个配合num_workers开太多有时候会额外占显存,尤其你2万张图读入时会有缓存堆积。另外ResNet50的BN层在迁移学习时如果开了梯度,反向传播也会多占一些,干脆把所有BN层冻住试试,能省不少显存。混合精度(AMP)确实立竿见影,显存几乎能砍一半,而且还快,建议你直接上torch.cuda.amp,对准确率基本没影响。梯度累积的话是曲线救国,相当于把大batch拆成小batch多次更新,训练慢的问题可以靠调高学习率或者用多步scheduler稍微补回来一点。还有个细节:你用的优化器是不是Adam?它和SGD比显存占用会多一截,可以换成SGD+momentum试试。最后,降batch size到8后准确率差可能不是显存的问题,而是学习率没跟着调,batch size变小了学习率也得对应降一点,不然收敛会不稳。
12G跑ResNet50加32的batch确实有点悬,224的图一张就得占不少显存,加上梯度回传和优化器状态,16到24之间其实更安全。你降到8能跑说明不是代码bug,优先试试梯度累积,设成4步等效batch32,显存压力小很多准确率也不会崩。混合精度也是个好办法,PyTorch自带amp包,基本一行代码改动就能省30%显存。另外num_workers设成4到6就够了,开太多反而容易挤占显存交换空间。