废话不多说直接进入正题。这一篇是这个系列的第四篇主题是很多人在语义分割学习中绕不过去的第一道坎——CPU和GPU版本的PyTorch及其环境安装。前几篇我们聊了语义分割的数据标注、模型结构、评价指标但真正动手跑代码的时候你会发现最先卡住你的往往不是网络结构而是环境。环境这玩意儿在语义分割项目里比在普通图像分类项目里更敏感数据增强库、图像读写库、GPU驱动、PyTorch版本任何一个环节松了训练就直接崩给你看。所以这一篇我把自己这么多次配环境的经验完整写出来包括为什么一个装法五花八门、哪些命令必须一条条敲、哪些坑踩一次就够。1. 语义分割学习里环境安装为什么值得单独写一篇1.1 环境是不可忽视的隐形变量语义分割项目看起来只是图像进去、mask出来但实际操作里它依赖的东西比想象中多得多。要读图、做数据增强、跑pre-trained backbone、计算逐像素的损失还要在训练中不断保存checkpoint和可视化结果。这一套流程里Python版本、torch和torchvision的版本、OpenCV的版本、CUDA运行时、以及显卡驱动全部被串联在一起。其中最典型的例子就是torch和torchvision的强绑定关系。很多人跑语义分割代码时报出ModuleNotFoundError: No module named torchvision.models...其实不一定是你代码写错了经常只是torchvision版本和你手头的torch版本不匹配。torch 2.1对应torchvision 0.16torch 2.2对应0.17torch 2.3对应0.18这个矩阵错一位都可能让你多折腾两个小时。而分割任务常用的那些pre-trained权重往往还依赖torchvision里某个指定版本的模型定义版本不对连参数都load不进去。所以在整个学习系列中间插这一篇是有意为之。前面几篇是让你理解分割任务在做什么这一篇是保证你后面的模型代码能真正跑起来而且每次复现实验结果时不会因为环境不同产生莫名其妙的差异。1.2 CPU与GPU的差别直接决定你的学习路径很多初学语义分割的同学会纠结一个问题我没有独显能不能学答案是能但你的学习路径和节奏要调整。CPU版的PyTorch在纯学习场景下完全够用。你可以用小数据集、小模型去理解分割网络的输入输出、理解损失函数怎么算、理解预测结果怎么转成可视化mask。不过一旦进入真实训练环节哪怕是训练一个很小的UNetCPU和GPU的差距也是数量级的。我自己在CPU上试过跑一个512×512输入的轻量分割模型一个epochCPU要跑接近二十分钟换到一张普通中端显卡只要一两分钟。这种差距不是忍一忍能解决的它会直接影响你调参的效率因为你根本等不起结果。更总要的是语义分割是所有视觉任务里对显存最不友好的方向之一。普通分类任务对一张图输出一个类别向量分割任务要输出和原图分辨率相当的逐像素预测显存占用成倍增长。同样的模型在分类任务里batch size可以开到32在分割任务里经常只敢开8甚至4。所以如果你决定用GPU训练有一个比有没有GPU更关键的问题你的显卡显存多大。6G显存和24G显存跑同一个模型能配置的策略完全不同。2. 动手安装前先把显卡驱动、CUDA和Python版本的关系理清楚2.1 驱动、CUDA Toolkit、cuDNN三者的分工与常见误解我第一次装GPU版PyTorch的时候犯过一个典型错误以为只要装了CUDA Toolkit就能跑GPU结果走了很多弯路。这里先把三者的关系说透。显卡驱动Driver是操作系统和GPU硬件之间的桥梁。没有它系统根本识别不了显卡nvidia-smi命令也会直接报错。CUDA Toolkit是一整套开发工具链里面包含编译器、运行时库、调试工具等。cuDNN则是NVIDIA专门为深度神经网络计算优化的加速库卷积等底层操作会调用它来提速。关键点是PyTorch的pip包和conda包本身已经捆绑了它运行所需的CUDA运行时库和cuDNN。也就是说你只要安装了GPU版的PyTorch它内部自带的库就能直接调用GPU计算并不需要你额外去装一套完整的CUDA Toolkit。真正需要手动装CUDA Toolkit的场景只有一个你要自己写或编译自定义的CUDA算子比如给某个分割网络写一个定制化的高效上采样模块那时才需要nvcc编译器。用一个不严谨但好理解的类比显卡驱动是电源插座CUDA Toolkit是全套工具箱cuDNN是工具箱里那把专用的电动螺丝刀。而PyTorch官方提供的GPU版安装包相当于厂家已经给你备好了电动螺丝刀你只需要保证电源插座能通电就行。那么这个通电条件是什么你的显卡驱动支持的CUDA版本必须不低于你装的PyTorch包所对应的CUDA版本。比如你想装cu118版PyTorch对应CUDA 11.8你的驱动必须支持CUDA 11.8及以上想装cu121版驱动版本要求就更高。驱动是向下兼容的新驱动能跑旧CUDA版本反过来则不行。2.2 三步确认自己的安装条件装环境之前强烈建议先把机器条件摸清楚。整个过程就三步花不了五分钟。第一步确认显卡驱动状态。在命令行输入nvidia-smi如果能正常输出一张包含显卡型号、驱动版本、显存信息的表格说明驱动是好的。请重点关注两个字段Driver Version驱动版本和CUDA Version这个数字代表当前驱动支持的最高CUDA版本并不是说系统里装了对应版本的CUDA Toolkit。如果这条命令报错比如提示nvidia-smi不是内部或外部命令那基本可以确定你的显卡驱动没装好或者是集成显卡环境根本没识别到独立显卡。第二步确认显卡的算力级别。这里的算力Compute Capability是一个硬件能力指标。NVIDIA官方每个GPU型号都有对应的算力编号比如老一点的GTX 10系算力在6.xRTX 20/30系在7.5和8.xRTX 40系在8.9。算力太低一般低于3.5的老显卡新版PyTorch不会再做支持装了也跑不起来。第三步确认Python版本。PyTorch对Python版本的支持有一定的滞后性太新的Python版本比如刚发布的大版本经常出现包还没编译好、安装失败的情况。做深度学习项目我比较推荐Python 3.9或3.10这两个版本兼容性最好社区生态里的所有库基本都能找到对应版本。3. CPU版PyTorch的安装流程与适用边界3.1 CPU版安装的具体操作CPU版安装其实是最简单、最不容易出错的因为它不涉及CUDA那一大堆问题。哪怕你以后有钱升级了显卡我也建议你先掌握CPU版的安装逻辑因为后面排查GPU问题时你会反复用到干净环境这个概念。第一步创建一个独立的虚拟环境。我个人的习惯是每个项目都单独建一个环境不要让Python全局环境变成一锅大杂烩。这里以conda为例conda create -n seg python3.9 conda activate seg第二步选择CPU版PyTorch的安装命令。最稳妥的方式是去PyTorch官网首页选择对应配置生成安装命令因为版本在持续推进我在这里给出一个示例命令但大家拿到文章的时候应该以官方生成的最新命令为准pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu这条命令里最关键的是--index-url参数。它指定了安装包的下载源。如果不写这个参数直接pip install torch你从默认PyPI源拿到的是什么版本就不一定了很可能是带CUDA依赖的体积很大的版本或者不是官方预编译包。在CPU环境下清楚地指定CPU版本能避免很多隐性问题。如果你坚持用conda方式装命令是这个结构conda install pytorch torchvision torchaudio cpuonly -c pytorch注意最后那个-c pytorch它表示从PyTorch官方conda频道安装。装完之后验证一下import torch print(torch.__version__) print(torch.cuda.is_available())CPU版的环境里torch.__version__不会带cu这种后缀torch.cuda.is_available()返回False这都是正常现象不要慌。3.2 CPU版能做什么不能做什么按我的经验CPU版在语义分割的学习阶段可以做这些事第一学习数据流。语义分割的完整流程里很大一部分其实和GPU无关。数据读取、Dataset类实现、预处理、mask的生成和可视化这些你用CPU版都能跑通而且跑通之后你对整个流程的理解会非常扎实。第二跑小规模的模型验证。如果你想快速验证某个网络结构的前向传播是否走通输入输出维度是否正确CPU版完全够。你可以把图像缩得很小、batch设成1快速跑一遍。第三用预训练模型做单张图片的推理。虽然慢一点但一张512×512的图推理一次CPU版也就是几十秒到几分钟的事完全可接受。但有两件事CPU版是真的扛不住。一是正经训练一个语义分割模型。分割任务本身计算量大再加上输入分辨率高CPU训练的效率低到会让你怀疑自己是不是代码写错了。二是混合精度训练。虽然CPU新版也支持了某些AMP操作但实际加速效果和GPU完全不是一个量级没必要折磨自己。所以我的建议很直接如果只是过一遍代码逻辑可以把CPU版当作调试环境但如果你认真要学语义分割该上GPU还是要上GPU。4. GPU版PyTorch的安装全流程4.1 pip方式安装GPU版PyTorchGPU版的安装核心就一句话选对你的CUDA版本然后一条命令装完。先回到前面nvidia-smi的输出那一行CUDA Version就是你的安装参考线。比如它显示CUDA Version: 12.1那你装CUDA 12.1版本对应的PyTorch包通常标记为cu121就没问题显示11.8就装cu118。原则是PyTorch包对应的小版本号不能高于驱动支持的CUDA版本号。具体安装命令我以cu121为例pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121和CPU版相比区别就是路径最后的cu121。这个路径是PyTorch官方专门预编译好的GPU版本wheel包目录里面已经捆绑了对应的CUDA运行时和cuDNN。我见过有些人装GPU版时写pip install torch然后发现torch.cuda.is_available()一直是False折腾半天也没搞明白为什么。这里要再次强调最好明确指定官方wheel源不要图省事。官方预编译包的命名规则一目了然带cu前缀的就是GPU版这能帮你从一开始就确认装的东西对不对。4.2 conda方式安装GPU版PyTorch如果你用的是conda环境那条对应的安装命令是这样conda install pytorch torchvision torchaudio pytorch-cuda12.1 -c pytorch -c nvidia和CPU版的区别在于多了一个pytorch-cuda12.1的包。conda会自动帮你把CUDA相关的库作为依赖装好这一点的确比pip省心不用自己操心额外依赖的问题。但conda的方式有一个需要警惕的点它会把nvidia频道纳入依赖解析如果你所在网络的下载速度一般这一步可能比pip要慢不少。关于pip和conda我的核心建议是选定一种方式并坚持到底。这两个包管理器的依赖解析逻辑不同混用时的冲突排查难度成倍上升尤其是torch这种重度依赖库pip装一半、conda装一半最后环境里出现两个torch版本互相打架的情况我见过太多次了。4.3 装完必须跑的GPU验证脚本让我把话放在前面如果你装完GPU版PyTorch不去验证那等于白装。下面的脚本是我每次配完环境必跑的直接复制粘贴就能用import torch print(PyTorch版本:, torch.__version__) print(CUDA是否可用:, torch.cuda.is_available()) if torch.cuda.is_available(): print(GPU数量:, torch.cuda.device_count()) print(GPU名称:, torch.cuda.get_device_name(0)) print(当前GPU设备索引:, torch.cuda.current_device()) # 跑一个简单的GPU矩阵乘法确认计算真的执行了 a torch.randn(1000, 1000).cuda() b torch.randn(1000, 1000).cuda() c torch.mm(a, b) print(GPU矩阵乘法验证通过输出形状:, c.shape)当你看到CUDA是否可用: True并且输出了显卡名称说明GPU版环境已经通了。这里顺便教大家一个鉴定版本的小技巧查看torch.__version__的输出内容里面如果包含类似cu121这样的后缀就说明你装的是带CUDA支持的GPU版如果只是2.x.x后面什么都没有那就是CPU版。跑这个脚本的同时你可以在另一个终端窗口输入nvidia-smi动态观察GPU的显存占用情况。如果你发现跑矩阵乘法时显存使用量往上跳了一截那说明PyTorch确实在调用GPU资源而不是只做了个表面文章。5. 语义分割项目跑起来之前周边环境也要一次配齐5.1 torchvision、OpenCV、数据增强库的版本搭配如果说PyTorch是语义分割的发动机那周边这些库就是变速箱、轮胎和方向盘缺一个都跑不动。首先是torchvision。我已经在前面强调过它和torch的强绑定关系这里再给一个通用的判断方法装完torch之后去PyTorch官网的版本矩阵页面查一下当前torch版本对应的torchvision版本严格对齐不要凭感觉装最新版。尤其是分割任务里经常用到torchvision提供的预训练backbone权重版本错了轻则警告重则加载失败。然后是OpenCV和图像处理相关的库pip install opencv-python tqdm matplotlib numpy pillow这里有个我踩过很多次的坑OpenCV和numpy的版本冲突。新版本的numpy比如1.24以上搭配老版本的opencv-python经常会出现module numpy has no attribute float之类的报错或者更隐蔽的numpy.ndarray size changed错误。解决方案不是自己去找所谓兼容版本而是全部装到较新版本numpy和opencv一起升级让它们处于同一代际。语义分割还强烈建议装一个库——albumentations。这个数据增强库对分割任务有一个巨大的优势它可以同时对输入图像和对应的mask做完全一样的随机变换而且你不需要自己关心坐标变换的映射关系。如果不用它你写水平翻转、随机裁剪、颜色抖动这些增强时图像和mask的同步处理很容易写错。5.2 训练脚本中与环境相关的参数设置环境装好之后训练脚本里还有几个跟硬件很相关的参数值得好好设一下而不是用默认值一跑到底。第一类参数是DataLoader的num_workers和pin_memory。num_workers控制数据读取的子进程数量语义分割的数据集往往包含大量高清图像和对应的maskCPU读图和解码本身是耗时操作如果num_workers0GPU利用率会非常难看。我一般从4起步再根据机器配置往上调。pin_memoryTrue能让数据在传输到GPU时省去一次拷贝在显存充足的条件下基本是必开的。第二类是设备管理代码。所有模型和tensor都必须在同一个设备上device torch.device(cuda if torch.cuda.is_available() else cpu) model model.to(device) images images.to(device) masks masks.to(device)这个逻辑看起来很简单但我见过无数分割项目的报错都出在这里模型在GPU上某个tensor忘记.to(device)然后一训练就报Expected all tensors to be on the same device。这里有个小习惯值得培养——写训练循环的时候统一把一个batch里的图像、mask、损失函数里的target全部放到同一个设备变量上不要这里一个一个写死cuda。第三类是混合精度的处理。GPU版本能带来的一个非常实用的加速手段就是AMPAutomatic Mixed Precision在语义分割中尤其有用因为分割模型通常较大半精度推理和训练可以显著降低显存压力。一个很简单的示范from torch.cuda.amp import autocast, GradScaler scaler GradScaler() with autocast(): outputs model(images) loss criterion(outputs, masks) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()需要注意的是混合精度虽然名字叫混合但它在CPU上几乎不会给你带来任何收益我还是把它归到GPU环境专属技能里面。6. 那些反复出现的环境坑从现象查到根因6.1 torch.cuda.is_available()返回False的完整排查链路这是环境配置里出现频率最高的疑难杂症。现象是驱动看起来是好的装的是GPU版PyTorch但一运行torch.cuda.is_available()就是False。别急着重装按下面这条链路一步步排查大多数情况十分钟内能定位。第一步重新确认你装的确实是GPU版。打开命令行输入python -c import torch; print(torch.__version__)如果输出的版本号后面没有cu字样问题就是明摆着的你装的是CPU版。原因可能是安装时走了默认PyPI源或者是conda/pip混用时不小心用另一个包管理器把torch覆盖掉了。解决方式是把环境里的torch先卸干净再用前面说的官方wheel源重装。第二步确认显卡能被系统识别。nvidia-smi如果正常显示说明驱动层面OK。如果这条命令直接报错那不用看PyTorch了先把驱动问题解决掉。不要试图用老版本PyTorch绕过驱动问题那只会引发更多莫名其妙的错误。第三步核对驱动的CUDA支持范围。nvidia-smi输出里的CUDA Version字段是你驱动支持的最高版本。如果你的PyTorch包是cu121版本但驱动支持的最高CUDA版本只有11.8那torch.cuda.is_available()返回False就非常合理了。这时候有两条路升级显卡驱动或者降低PyTorch的CUDA版本。第四步检查是否同一个环境里同时存在多个torch副本。这个问题在conda环境里尤其常见conda list显示一个torchpip list又显示另一个torch。两套包管理系统各自为政最后import的到底是哪个完全取决于环境路径顺序。最干净的处理办法是把环境删掉重建这次老老实实用同一种安装方式。6.2 显存不足与训练速度异常的排查思路GPU环境跑通之后第二个经常把人卡住的硬伤就是显存爆掉。语义分割训练时CUDA out of memory基本是每个人都会遇到一次。显存不足的第一反应是调小batch size。我从8调到4还不行从4调到2再不行就调到1。这当然是最粗暴有效的办法但不应该一上来就这么做。你应该先判断是不是代码里有些变量没有及时释放比如在循环里反复创建历史tensor、没有用del释放中间变量、验证阶段也没有关掉梯度计算等等。这些释放问题解决了可能batch size根本不用降。如果确定是模型本身超显存了可以组合使用几招一是缩小图像输入尺寸分割任务普遍从512或者768起步降到384甚至256显存压力会显著缓解代价是精度会有可感知的下降二是用混合精度半精度tensor的显存占用直接砍半变少三是用梯度累积模拟更大的batch size比如显存只够batch 4你想用batch 16的效果就每个batch正常反传并累积4步梯度后更新一次参数。再说训练速度异常的问题。如果你发现GPU利用率很低显存占用上去了但利用率总在20%到30%徘徊大概率瓶颈在数据加载。CPU读图、resize、做数据增强的速度跟不上GPU计算速度GPU空转等待。这时候你需要做的是提高num_workers让多个子进程并行读数据对数据增强做一个权衡训练阶段不要堆叠太多耗时操作另外检查一下是不是在CPU上额外做了大量预处理运算。还有一种速度异常是混合精度没开。分割模型普遍又大又深FP32训练比FP16训练在相同显卡上的速度差距是肉眼可见的。如果你确认GPU环境没问题但训练速度远低于别人在类似显卡上的报告数字建议先检查脚本里有没有使用AMP。收尾的一点个人经验环境安装这一关说难也不难说简单也不简单。我把折腾过的这些细节写出来核心就是想传达一件事先确认条件再动手安装最后做验证一步一步来不要相信一条命令搞定所有环境的神话。我的习惯是每配好一套环境就把当前的torch版本、CUDA版本、显卡驱动版本、周边库版本记录下来存成一个文本文件放在项目目录里。以后再遇到问题或在别的机器上复现实验这份环境记录能帮我省掉大量无意义的排查。还有一个小技巧值得分享虚拟环境的名字也能帮你记忆我一般把环境直接命名为类似seg_cu121_py39的格式一眼就能看出这个环境里装了什么版本的CUDA和Python。这一篇把地基打牢了下一篇我们就能正式进入语义分割项目的数据准备和训练实战。到时候你会发现真正跑起训练来环境这些事先做好的准备才是让你能把精力全部聚焦在模型与实验本身的最大底气。