资讯动态

Windows PyTorch训练ResNet-50 ImageNet-1K避坑

发布时间:2026/9/30 10:47:28 来源:尧图企业网站定制
1. 在Windows上跑ResNet训练这件事先想清楚值不值得很多人一提到ImageNet-1K级别的训练脑子里第一反应就是这活儿得在Linux服务器上干理由也充分多进程DataLoader在Linux上是fork开销小文件系统IO调度成熟工具链和论文复现脚本默认都在Linux环境里验证过。但现实是相当一部分做CV方向的同学、工程师手头的主力机器就是一台装好显卡的Windows工作站或者带独显的笔记本要额外配一台Linux机器成本、维护精力、跨机器传数据的麻烦都不划算。所以我这几年一直坚持把训练链路搭在Windows上从环境配置到数据读取到训练循环每个环节都自己踩一遍。这篇文章讲的就是在Windows下用PyTorch把ResNet在ImageNet-1K上训练起来的完整经验包含可直接复用的代码、环境版本对应关系、以及Windows平台独有的一堆坑。它解决的核心问题不是ResNet怎么定义这种教科书问题而是为什么同样的脚本在Linux上跑得好好的搬到Windows就慢、就报错、就卡死。适合谁看已经会写基本PyTorch训练脚本、想用自己机器跑中等规模图像分类的人或者已经在Windows上装好环境但训练速度远低于预期的人。如果你连conda和显卡驱动都还没装这篇文章也能当一份按图索骥的安装清单只是进度会稍微慢一点。我先把结论摆出来Windows跑ImageNet训练完全可行单机单卡ResNet-50一个epoch128万张图在合理配置下是可以压缩到可用时间区间的关键是数据处理流水线要调对而这恰恰是Windows上最容易出问题的地方。1.1 为什么很多人默认用Linux而我还是选了Windows先承认Linux的优势是客观存在的。DataLoader的多进程在Linux上有fork可用子进程能直接继承父进程的内存映射启动成本低Windows从Python层面只能用spawn方式启动子进程每个worker都要重新import一遍主模块重新初始化一遍环境这就是为什么在Windows上把num_workers调大反而可能变慢甚至内存爆炸。这个差异是真实的不是心理作用。但我选Windows的理由也很实在。第一我的数据集和标注工具、可视化脚本、后续的推理部署验证都在Windows上数据在本地NVMe盘上读起来带宽够。第二Windows现在的WDDM驱动模型虽然有一点调度开销但在单卡训练场景下这点开销相比数据加载的瓶颈几乎可以忽略。第三也是最重要的一点我在一次次的踩坑过程中把Windows上那几个固定的坑摸清楚了一旦跨过去稳定性并不比Linux差。很多Linux一定更快的说法其实是拿一个调优过的Linux环境和一个没调优的Windows环境作对比不公平。所以我的态度是环境选哪个不重要重要的是你知不知道这个平台的边界在哪。下面我会把Windows上那几个真实存在的边界一个一个拆开讲。1.2 Windows和Linux在训练链路上的三处真实差异第一个差异是进程启动模型。Linux是forkWindows是spawn。spawn意味着每个DataLoader worker都会重新执行你的主脚本的顶层代码所以如果你把模型定义、数据集实例化、甚至训练循环都写在模块顶层而不放在ifname main:保护块里Windows下就会无限递归创建进程直接报错或者把内存吃光。这是我第一次在Windows上跑训练时最先撞上的墙。第二个差异是共享内存。Windows没有Linux那种/dev/shmPyTorch的DataLoader在传递张量时用的共享内存机制在Windows上走的是另一套实现当num_workers开得比较大、每个batch又比较大的时候容易碰到共享内存不足或者worker进程莫名退出的问题。表现就是训练中途某个epoch突然卡住日志不再刷新。第三个差异是文件句柄和路径。Windows对同时打开的文件句柄数量、路径长度历史上260字符限制、以及反斜杠转义都比较敏感。ImageNet的目录是train/n01440764/xxx.JPEG这种结构路径不长问题不大但如果你把数据集放在很深的嵌套目录里或者目录名带空格和中文就会时不时冒出读取失败。这三处差异后面我会用专门的章节给对应的修法。先把环境搭起来。2. 环境链路从显卡驱动到能import torch的完整搭建Windows下装PyTorch出问题的从来不是pip install本身而是版本对应关系。我会按显卡驱动 → CUDA运行时 → PyTorch → torchvision这条链路讲每一环都要能对上下一个。你可能听过CUDA Toolkit必须装其实用官方的pip wheel安装PyTorch时wheel里已经自带了对CUDA运行时库的依赖你不一定需要在系统里单独装完整的CUDA Toolkit只需要显卡驱动版本足够新能支持对应的CUDA版本就行。但如果你后续要自己编译CUDA扩展那就得装Toolkit这是两回事。我建议的路线是先确定显卡驱动支持的CUDA上限再选定PyTorch版本最后用conda或者pip安装对应cu版本的wheel。整个过程不需要手动下载CUDA安装包除非你有编译需求。2.1 CUDA、驱动、PyTorch三者版本怎么对上这块的逻辑很多人搞混。显卡驱动里包含了一个驱动能支持的最高CUDA运行时版本这个版本号可以在命令行用nvidia-smi查看右上角会显示CUDA Version: xx.x。注意这个数字不是你装的CUDA版本而是驱动能兼容的最高版本。你完全可以在驱动支持12.x的情况下安装cu118CUDA 11.8版本的PyTorch这是向下兼容的没问题。反过来就不行如果你的驱动只支持到11.7你却装了需要CUDA 12.x运行时的PyTorch wheel那import torch的时候就会报错或者torch.cuda.is_available()返回False。所以顺序是先看nvidia-smi的版本号再决定装哪个cu前缀的PyTorch。下面这张表是我实际用过的几组稳定搭配供参考显卡驱动CUDA上限建议PyTorch版本wheel的CUDA前缀备注12.4及以上2.2 / 2.3cu121目前最省心的组合12.0 - 12.32.1 / 2.2cu118兼容性广够用11.7 - 11.81.13 / 2.0cu117老机器上常见11.6及以下1.12cu116建议先升级驱动有人会问conda装和pip装有什么区别。conda装的PyTorch会把cudatoolkit作为conda包一起装上好处是它自己在conda环境里维护一份CUDA库不会和系统冲突pip装的PyTorch则是依赖wheel里打包的运行库。两者都能用我个人偏向conda建环境、pip装torch的组合因为conda对Windows上一些C依赖比如微软的VC运行时处理得更省心。2.2 conda环境与依赖清单我习惯给每个项目单独建环境避免不同项目的torch版本互相污染。命令很简单conda create -n resnet_train python3.10 -y conda activate resnet_trainPython选3.10是我这几年最稳的3.11、3.12在新一点的torch上也行但一些周边的库比如早期版本的tensorboard、一些编译型依赖在3.12上偶尔没有预编译wheel会触发源码编译在Windows上编译是很痛苦的。所以保守起见3.10能省很多事。然后是核心依赖pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121 pip install numpy pillow tqdm tensorboard这里用--index-url指向官方的CUDA wheel仓库能确保拿到的是cu121版本而不是默认的CPU版本。我踩过的坑就是直接pip install torch结果装成了CPU版训练跑起来torch.cuda.is_available()是False还以为是驱动坏了排查半天。装完之后可以顺手确认一下python -c import torch; print(torch.__version__, torch.version.cuda, torch.cuda.is_available())输出形如2.2.0 12.1 True就对了。如果最后是False先别急着重装按下面的顺序查。2.3 验证环境是否真的能用到GPU排查is_available为False我一般的顺序是第一确认nvidia-smi能正常输出如果这个命令本身都找不到或者报错那是驱动没装好先去装显卡驱动这一步和PyTorch无关第二确认装的是CUDA版wheel用上面那条命令看torch.version.cuda是不是None如果是None说明装成了CPU版重装第三确认驱动版本上限不低于wheel的CUDA版本如果wheel是cu121而驱动上限只有11.7就降级装cu118。再补一个硬件层面的真实检查import torch print(设备数:, torch.cuda.device_count()) print(设备名:, torch.cuda.get_device_name(0)) x torch.randn(1024, 1024, devicecuda) y x x torch.cuda.synchronize() print(矩阵乘完成, 结果均值:, y.mean().item())前面叫环境能识别GPU这段叫环境能真正算东西。有些情况下识别没问题但一跑算子就报非法内存访问通常是驱动和运行时之间的兼容问题这时候更新驱动是最直接的解法。3. ImageNet-1K数据集的落地目录结构、预处理与Windows读取瓶颈环境好了接下来是最耗精力的一环数据。ImageNet-1K是1000类的分类数据集训练集大约128万张验证集5万张全部是JPEG格式的图片尺寸不一基本都在几百像素这个量级。它的目录组织方式是固定的train目录下按类别ID分子目录每个子目录里放该类的图片val目录平铺所有图片另外配一个验证集标注文件来告诉程序每张图属于哪一类。理解这个结构很重要因为它是决定你怎么写Dataset和怎么配transforms的前提。这里我要强调一个很多人忽略的点数据放在哪种盘上对Windows训练速度的影响比换GPU还大。我用过机械硬盘、SATA SSD和NVMe SSD三种同样的ResNet-50脚本仅数据加载这一块的时间差就能拉到三四倍。原因是ImageNet的图都是小文件随机读取密集机械盘的随机IO完全扛不住。所以如果你的数据还在机械盘上先把它挪到SSD这一步的收益立竿见影比调任何参数都来得实在。3.1 数据集结构必须长成什么样标准结构长这样imagenet/ train/ n01440764/ n01440764_10026.JPEG ... n01443537/ ...共1000个目录 val/ ILSVRC2012_val_00000001.JPEG ... val_label.txtval_label.txt的每一行是文件名 类别索引因为是平铺结构torchvision的ImageFolder处理不了需要自己写一个Dataset。我下面的实现会把这个逻辑写清楚。有些人找的val分好类的版本其实是把验证集也按类别放了子目录那种版本可以直接用ImageFolder但要注意类别顺序必须和训练集一致否则算出来的准确率是错位的。判断方法很简单看一眼train的1000个目录名排序和val的目录名排序是否完全一样一样就可以放心用ImageFolder。我实际用的写法是自定义Dataset读val_label.txt更可控。代码在后半部分给。3.2 预处理流水线里每个参数的道理训练用的transforms我固定用这一套它是ImageNet上的事实标准from torchvision import transforms train_tf transforms.Compose([ transforms.RandomResizedCrop(224, scale(0.08, 1.0)), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]), ])逐个说为什么。RandomResizedCrop(224, scale(0.08, 1.0))是随机裁剪一块面积占比8%到100%的区域再缩放到224这是数据增强的核心让模型对物体位置和尺度不敏感scale下限0.08看起来很小但它保证了裁出来的小块至少还有一定内容是原论文验证过的稳妥值。RandomHorizontalFlip是随机水平翻转对自然图像来说翻转不改变语义能免费增大数据多样性但在字符识别或医学影像这类任务里要谨慎。Normalize里的mean和std是ImageNet统计出来的像素均值方差不是随便填的。这三个通道的数值来自对整个训练集RGB通道的统计如果你不做这一步归一化训练初期很容易出现梯度爆炸或者loss震荡。这个坑我见得太多有人训练loss一直下不去换成这套标准化立刻就好了。验证集的处理要克制只做确定性操作val_tf transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]), ])先缩到短边256再中心裁224这是标准的评估流程。用随机增强去做验证会让每轮的val准确率波动很大无法客观反映模型状态。我刚入门时就犯过这个错把增强也用到验证上val准确率忽高忽低误以为模型不稳定。3.3 Windows下数据读取慢的根因与缓解前面说了Windows的进程启动是spawn这在数据加载上的直接后果是num_workers每加一个就多一次Python解释器冷启动、多一次torch库的加载、多一次数据集路径扫描ImageFolder会扫描全部1000个目录这个扫描在128万张图上是要花时间的。所以num_workers不是越大越好。我实测下来8核到16核的CPUnum_workers给4到8比较合适。给到16反而会因为进程调度和内存复制开销速度不升反降。而且要注意spawn方式下每个worker都会独立执行一遍你脚本顶层的代码如果你在顶层实例化了ImageFolder那就是每个worker都扫一遍全部目录非常浪费。正确做法是把数据集实例化放在ifname main:内部只执行一次。缓解读取慢的第二个手段是关掉多余的图像校验。PIL加载图片时默认会做一些检查大流量读取时这些检查也是开销。可以用下面这种写法跳过from PIL import Image, ImageFile ImageFile.LOAD_TRUNCATED_IMAGES True把LOAD_TRUNCATED_IMAGES设为True可以避免因为个别损坏的JPEG导致整个读取报错中断这在ImageNet上是很有必要的因为样本量大难免有几张图是坏的。但要注意这个设置在Linux上也是一样的道理只是Windows上文件系统对文件锁更敏感遇到损坏文件时更容易卡住。第三个手段是把预处理的一部分下沉到worker里也就是把ToTensor和Normalize放在Dataset的__getitem__里而不是在训练循环里做。这样每个worker并行地完成CPU密集的转换GPU只管前向反向。这个原则在Linux和Windows上通用但在Windows上尤其重要因为一旦CPU成了瓶颈spawn的额外开销会让整体更慢。4. ResNet训练脚本的关键组成与代码实现到这一节环境、数据都到位了我来把这个训练脚本拆开讲。我会给一份完整的、能在Windows单卡上直接跑的代码同时解释每一处为什么这么写。这份代码不是最花哨的但每一行我都在真实数据上验证过。为便于阅读我拆成几个部分讲模型构建、优化器与调度、训练循环、混合精度。这里的数据集路径、batch大小这类参数都会给出我实测过的取值。先说一个原则训练脚本的所有顶层执行代码都必须放在ifname main:保护块内。这不是风格问题是Windows下能不能跑起来的硬性要求。4.1 模型构建与是否加载预训练权重关于要不要用预训练权重我的态度是分情况。如果你是在做研究、想复现ImageNet上的基准精度从零训练是必要的但如果你是在自己的业务任务上做迁移学习哪怕数据量在百万级加载ImageNet预训练权重再微调收敛速度和最终精度都明显更优。很多人以为数据够多就不用预训练实际上预训练省的不只是时间还有对超参的敏感性。模型直接调torchvision的实现最省事但要改最后一层适配类别数虽然ImageNet就是1000类不改也行import torch.nn as nn from torchvision import models def build_model(num_classes1000, pretrainedFalse): model models.resnet50(weightsNone) # 从零训练 # 若加载预训练models.resnet50(weightsmodels.ResNet50_Weights.IMAGENET1K_V1) if num_classes ! 1000: model.fc nn.Linear(model.fc.in_features, num_classes) return model这里有个版本提醒老版本torchvision用pretrainedTrue参数新版本改成了weights枚举两者混用会报错或者给出警告。我建议统一用weights参数明确写清楚用哪个预训练版本避免以后换环境时出问题。另外如果你在Windows上多卡并行要用nn.DataParallel或者DistributedDataParallel。单卡就直接.cuda()。多卡场景下Windows的DDP配置比Linux麻烦一些涉及端口、NCCLWindows上其实走的是Gloo或新的后端这个如果环境不是多卡建议先不碰把单卡跑通再说。4.2 优化器、学习率调度与损失ResNet的经典训练配比是SGD momentum weight decay加上warmup和余弦退火。具体参数optimizer torch.optim.SGD( model.parameters(), lr0.1, momentum0.9, weight_decay1e-4 )学习率0.1是配合batch_size256用的标准值。如果你显存小只能开batch_size64那学习率要按线性缩放大致降到0.025不然容易训崩。这个线性缩放法则在小批量上不是绝对精确但作为起点非常好用。momentum 0.9和weight decay 1e-4也是ImageNet上的经验值改动它们之前先确认你有充分的理由。调度器我用余弦退火它能平滑地降低学习率避免手动调stepfrom torch.optim.lr_scheduler import CosineAnnealingLR scheduler CosineAnnealingLR(optimizer, T_maxepochs, eta_min1e-6)如果你的训练轮数少比如只训30轮余弦退火能把最后阶段的学习率压下去对收敛有帮助。如果轮数很多90轮以上可以再加warmup前5个epoch把学习率从很小线性升到0.1避免一开始梯度太大这在从零训练时尤其重要。损失函数用交叉熵注意PyTorch的CrossEntropyLoss内部已经包含了softmax你的模型最后一层不要自己再加softmaxcriterion nn.CrossEntropyLoss()我见过有人在模型里加了softmax外面又用CrossEntropyLoss结果loss一直不下降排查很久才发现在这里重复了。4.3 混合精度与显存控制Windows单卡显存通常有限8G、12G、24G不等开启混合精度AMP能显著降低显存占用并提升速度。原理是用FP16做前向反向计算同时保留一份FP32的主权重并用动态损失缩放来防止FP16下梯度下溢。用法很简单from torch.cuda.amp import autocast, GradScaler scaler GradScaler() with autocast(): output model(images) loss criterion(output, targets) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() optimizer.zero_grad(set_to_noneTrue)四个步骤一个都不能少。scale是把loss放大防止小梯度在FP16里变成0backward在放大后的loss上做scaler.step是先把梯度还原再更新权重如果这一步检测到梯度里有inf或nan会跳过这次更新scaler.update是动态调整缩放系数。很多人只写了前面两步忘了后面的step和update结果就是混合精度下训练不动。开AMP之后ResNet-50的显存占用能降大约三到四成速度能有百分之二三十的提升收益非常可观。但要注意AMP不是万能的如果你的模型里有大量小数值运算比如某些归一化可能需要谨慎处理。4.4 完整训练循环代码把上面的拼起来给出完整的可运行版本。注意每个函数和实例化都在main保护块内import os import time import torch import torch.nn as nn from torch.utils.data import DataLoader, Dataset from torchvision import transforms, models from torch.optim.lr_scheduler import CosineAnnealingLR from torch.cuda.amp import autocast, GradScaler from PIL import Image, ImageFile from tqdm import tqdm ImageFile.LOAD_TRUNCATED_IMAGES True class TrainFolder(Dataset): ImageNet train目录train/类别ID/图片 def __init__(self, root, transform): self.transform transform self.samples [] self.classes sorted(os.listdir(root)) self.class_to_idx {c: i for i, c in enumerate(self.classes)} for c in self.classes: cdir os.path.join(root, c) for f in os.listdir(cdir): self.samples.append((os.path.join(cdir, f), self.class_to_idx[c])) def __len__(self): return len(self.samples) def __getitem__(self, idx): path, label self.samples[idx] img Image.open(path).convert(RGB) return self.transform(img), label class ValFolder(Dataset): ImageNet val平铺图片 val_label.txt def __init__(self, root, label_file, transform): self.transform transform self.root root self.samples [] with open(label_file, r) as fp: for line in fp: name, cls line.strip().split() self.samples.append((name, int(cls))) def __len__(self): return len(self.samples) def __getitem__(self, idx): name, label self.samples[idx] img Image.open(os.path.join(self.root, name)).convert(RGB) return self.transform(img), label def accuracy(output, target, topk(1, 5)): with torch.no_grad(): maxk max(topk) _, pred output.topk(maxk, 1, True, True) pred pred.t() correct pred.eq(target.view(1, -1).expand_as(pred)) res [] for k in topk: correct_k correct[:k].reshape(-1).float().sum(0) res.append(correct_k.mul_(100.0 / target.size(0))) return res def main(): device torch.device(cuda if torch.cuda.is_available() else cpu) epochs 60 batch_size 256 workers 8 train_tf transforms.Compose([ transforms.RandomResizedCrop(224, scale(0.08, 1.0)), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]) val_tf transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]) train_set TrainFolder(rD:\imagenet\train, train_tf) val_set ValFolder(rD:\imagenet\val, rD:\imagenet\val_label.txt, val_tf) train_loader DataLoader( train_set, batch_sizebatch_size, shuffleTrue, num_workersworkers, pin_memoryTrue, persistent_workersTrue, prefetch_factor2, ) val_loader DataLoader( val_set, batch_sizebatch_size, shuffleFalse, num_workersworkers, pin_memoryTrue, persistent_workersTrue, ) model models.resnet50(weightsNone).to(device) if torch.cuda.device_count() 1: model nn.DataParallel(model) criterion nn.CrossEntropyLoss() optimizer torch.optim.SGD(model.parameters(), lr0.1, momentum0.9, weight_decay1e-4) scheduler CosineAnnealingLR(optimizer, T_maxepochs, eta_min1e-6) scaler GradScaler() for epoch in range(epochs): model.train() t0 time.time() running_loss 0.0 for images, targets in tqdm(train_loader, descfEpoch {epoch1}): images images.to(device, non_blockingTrue) targets targets.to(device, non_blockingTrue) optimizer.zero_grad(set_to_noneTrue) with autocast(): output model(images) loss criterion(output, targets) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() running_loss loss.item() scheduler.step() model.eval() top1, top5 0.0, 0.0 with torch.no_grad(): for images, targets in val_loader: images images.to(device, non_blockingTrue) targets targets.to(device, non_blockingTrue) with autocast(): output model(images) a1, a5 accuracy(output.float(), targets, topk(1, 5)) top1 a1.item() top5 a5.item() n len(val_loader) print(fEpoch {epoch1} | loss {running_loss/len(train_loader):.4f} f| top1 {top1/n:.2f} | top5 {top5/n:.2f} f| {time.time()-t0:.1f}s) torch.save({ epoch: epoch, model: model.state_dict(), optimizer: optimizer.state_dict(), scheduler: scheduler.state_dict(), }, fresnet50_epoch{epoch1}.pth) if __name__ __main__: main()这份代码可以直接跑把数据路径换成你自己的就行。下面几节我会讲它是怎么一步步踩坑调出来的。5. Windows下DataLoader多进程与那些一定会踩的坑这一节是全文最值钱的部分因为上面那份代码里的某些参数是我被反复折磨之后才加上的。Windows下的DataLoader和Linux差别很大很多在Linux上不出问题的写法在这里会以各种意想不到的方式炸掉。我把最典型的几类问题按处理顺序讲。5.1 num_workers配合main保护第一个坑也是最常见的忘记ifname main:或者只保护了main()调用但数据集和模型的实例化放在了顶层。Windows上spawn的worker会在启动时重新执行整个模块如果你的顶层有实例化代码那么每个worker启动时都会重新跑一遍轻则浪费几分钟扫目录重则因为某些资源比如端口、文件句柄被重复占用而直接报错。正确的结构是所有会执行的操作包括建数据集、建模型、建loader、开训练全部放在main函数里然后在模块底部用ifname main: main()来启动。类定义、函数定义可以放顶层因为它们只是定义不会执行。我带人做项目时看一份Windows训练脚本能不能跑第一眼就看这个结构对不对。5.2 共享内存不足与句柄泄漏第二个坑是worker进程在训练中途静默退出。表现是某个epoch跑到一半进度条不动了CPU占用掉下去GPU利用率掉到个位数日志也不刷新。这个原因通常有两个共享内存不够或者句柄泄漏。Windows没有Linux那种dev/shmPyTorch在传递张量时依赖系统的共享内存段每个batch的数据复制都要走这块。当你num_workers开得大、batch也大的时候共享内存段会迅速膨胀超过系统上限后worker就崩了。解决办法有这么几个降低num_workers我一般从8往下试到4降低prefetch_factor默认是2可以设成1甚至不设或者减小batch_size。句柄泄漏则是另一个方向的坑每个worker打开的文件如果没正确关闭累积到一定程度就报too many open files。PIL的Image.open默认是惰性加载的如果不调用load()或者不关闭句柄会一直挂着。稳妥的写法是在__getitem__里用with语句或者显式地img.load()之后让对象被回收。我在数据量大的时候会在__getitem__末尾加一句img.close()虽然有点啰嗦但能避免长跑训练时句柄慢慢堆积。另外persistent_workersTrue这个参数值得说道。它让worker在epoch之间不销毁、复用好处是避免了每个epoch重新spawn的开销。但代价是它要求数据集比较稳定如果你在训练中途动态改数据集内容可能拿到过期数据。对固定的ImageNet来说开它是划算的。5.3 常见报错对照表这几年我记录下来的Windows训练报错整理成一张表方便你按现象查报错或现象大概率原因处理办法RuntimeError: An attempt has been made to start a new process...缺少main保护把执行代码放进ifname main:DataLoader worker (pid xxx) is killed共享内存/内存不足降num_workers降prefetch_factorBrokenPipeError / EOFErrorworker崩溃导致管道断开先修worker崩溃的根因OSError: [Errno 24] Too many open files文件句柄泄漏getitem里显式close或withCUDA out of memory显存不足开AMP、降batch、梯度累积loss长时间为nan学习率过高或精度问题降lr、检查是否重复softmax、调AMP验证准确率乱跳验证用了数据增强验证只做ResizeCenterCrop训练速度远低于预期数据在机械盘 / worker过多换SSDnum_workers调4-8这张表可以贴在电脑边上出问题先对号入座能省掉大量时间。我要特别说BrokenPipeError它本身不是根因是worker已经崩了之后主进程试图和它通信的结果。很多人看到这个报错就四处改管道相关的代码其实应该往上翻日志找到worker真正崩的那一行。6. 训练过程中的调优与显存优化实战环境和数据都通了之后真正的折磨才刚开始为什么我的显存不够、为什么收敛这么慢、为什么精度上不去。这一节我讲几个实战中调出来的经验。6.1 学习率与warmup的真实作用学习率是训练里最敏感的参数。前面说0.1配batch 256这个值在从零训练ResNet时是标准配置。但它是配了warmup才稳的如果直接上0.1而你的权重是随机初始化的前几百步梯度会非常大loss很容易直接nan。warmup就是让学习率在前几个epoch从接近0线性升到0.1给网络一个适应期。warmup的实现可以用LambdaLR写from torch.optim.lr_scheduler import LambdaLR warmup_epochs 5 def lr_lambda(epoch): if epoch warmup_epochs: return (epoch 1) / warmup_epochs return 0.5 * (1 __import__(math).cos( (epoch - warmup_epochs) / (epochs - warmup_epochs) * 3.1415926)) scheduler LambdaLR(optimizer, lr_lambda)这段把warmup和余弦退火合并在一起了。实测下来从零训练加warmup前几个epoch的loss曲线会平滑很多不会再出现那种突然炸掉然后需要重启的情况。如果你是从预训练权重开始微调warmup可以短一些甚至不要因为初始权重已经不是随机的了学习率可以稍微小一点比如0.01到0.05之间具体看你的数据规模和任务相似度。6.2 梯度累积模拟更大batch显存不够是常态。假设你只有12G显存batch 256开不起来但你又不想改学习率因为改了要重新调这时候梯度累积是好办法。原理是一次前向反向算小批量的梯度但不立即更新攒够若干个步数再更新一次。这样等效于大batch但显存只按小batch算。代码上是把前面的训练循环改一下把scaler.step和zero_grad按累积步数做条件accum_steps 4 for i, (images, targets) in enumerate(train_loader): images images.to(device, non_blockingTrue) targets targets.to(device, non_blockingTrue) with autocast(): output model(images) loss criterion(output, targets) / accum_steps scaler.scale(loss).backward() if (i 1) % accum_steps 0: scaler.step(optimizer) scaler.update() optimizer.zero_grad(set_to_noneTrue)注意loss要除以accum_steps这样累加的梯度和一次性大batch的梯度在量级上是一致的否则梯度会成倍放大。这个细节很多人漏掉导致训练发散。我自己有次就忘了除loss直接起飞排查了好一阵。另外要小心训练集长度不能被accum_steps整除时最后一批处理不干净一般可以忽略或者手动补一次step。6.3 显存不够时按顺序尝试的几种手段显存不足时我会按这个顺序试第一开AMP混合精度收益最大通常能省三成显存。第二降batch_size这个是直接办法配合梯度累积保持等效批量。第三用channels_last内存格式把模型和输入都转成通道在后的格式这对卷积网络有时能提速省显存写法是model model.to(memory_formattorch.channels_last)输入也做同样转换。第四如果还不行考虑梯度检查点gradient checkpointing用时间换空间把中间激活值不保存反向时重算写法上对ResNet需要改结构比较繁琐作为最后手段。这里给个我实测的参考数据ResNet-50、224输入、batch 256不做任何优化时大约需要22G左右显存开了AMP能降到15G上下再配合梯度累积4步batch 64就能在8G显存的卡上跑起来。这个数字会随具体环境浮动但量级上是这样。7. 验证、断点续训与日志监控训练跑起来之后你需要能观察它、能中断它、能从中间恢复它。这一节讲怎么把这三件事做扎实。7.1 Top-1与Top-5验证实现ImageNet的标准评估是Top-1和Top-5准确率前者看最高分是否命中后者看前五是否命中。实现我写在4.4节的accuracy函数里了核心是取topk、然后比较预测和目标。这里有个容易错的地方验证时要model.eval()并用torch.no_grad()否则BatchNorm的统计量会被验证数据污染而且会白白占用显存。还有预测结果要转回FP32再算topk因为AMP下output是FP16直接算topk在极端情况下会有精度问题所以我代码里写了output.float()。验证集的batch不用和训练一样大它可以开大一点因为只做前向显存压力小。但val的num_workers也别开太大验证本身快worker太多反而在进程管理上浪费时间。7.2 checkpoint保存与断点续训长跑训练最怕的是跑到一半断电或者系统更新强制重启。所以每个epoch存一次checkpoint是必须的。我保存的内容包括模型权重、优化器状态、调度器状态和当前epoch号因为断点续训时优化器里的动量和调度器的学习率状态如果丢了恢复后的曲线会和中断前对不上。恢复的写法def load_checkpoint(path, model, optimizer, scheduler): ckpt torch.load(path, map_locationcpu) model.load_state_dict(ckpt[model]) optimizer.load_state_dict(ckpt[optimizer]) scheduler.load_state_dict(ckpt[scheduler]) return ckpt[epoch] 1map_locationcpu很重要它保证在CPU上先把权重读进来再搬到GPU避免直接往已经占满的显存里塞导致OOM。这个技巧在很多恢复场景都用得上。日志方面我一般除了print还会把每轮的loss、top1、top5、耗时写进一个CSV方便画曲线。tensorboard也可以但Windows上tensorboard有时候会因为端口占用或者临时目录权限问题启动失败纯CSV最稳配合pandas随手就能画图。8. 一些不成体系的碎碎念写到这里训练链路基本完整了。我最后说几个不那么技术、但确实影响体验的点。第一Windows的自动更新和杀毒软件是训练的两大隐形杀手。系统更新会强制重启把没存checkpoint的训练干掉杀毒软件实时扫描会在你读数据集的时候疯狂占用IO我遇到过一次训练速度莫名掉一半排查半天是杀毒在扫JPEG。建议训练前关掉自动重启把数据集目录加入杀毒白名单。第二GPU温度要盯着。Windows上跑满载训练笔记本或者散热不佳的台式机容易过热降频表现为训练前期快、后期慢。我一般用nvidia-smi每隔几秒看一眼温度和功耗确认没有降频。第三多卡训练在Windows上不是不能做但配置成本比Linux高。如果你的模型是单机能塞下的先把单卡调通、调好再考虑多卡。多卡不是线性的有时候两张卡的利用率加起来还不如一张卡调得好。最后分享一个我自己的小习惯我会在训练脚本里加一个--smoke参数跑一个极小的子集比如每个类抽10张图总共1万张用来在正式启动前快速验证环境、数据路径、前向反向、checkpoint保存这一整条链路是否通畅。这个习惯帮我省掉了无数次启动之后等两小时才发现路径写错了的尴尬。链路通了再放全量数据心里踏实。

读完文章,也想定制专属网站?

尧图设计师 24 小时内与您沟通定制方案

免费获取报价 →
↑