ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

从PyTorch入门到Unet图像分割:环境搭建与模型复现全攻略

从PyTorch入门到Unet图像分割:环境搭建与模型复现全攻略 简介面向深度学习入门与进阶群体的PyTorch学习代码与模型复现资源适合正在学习神经网络实现、想通过完整代码理解Unet、Vision Transformer等经典架构的读者。资源整理自龙良曲老师的PyTorch教程并在其基础上增加了多类模型的复现实验覆盖图像分割、视觉Transformer、无监督自编码器以及扩散模型DDPM等方向可帮助学习者将理论快速落实到可运行代码中。压缩包共259个文件约18.88MB以145个Python源码文件为核心辅以25张结构图/效果图、15个配置文件、7个Markdown说明文档及若干Jupyter Notebook示例层次清楚便于对照源码查看训练与可视化结果。此外Notebook示例与GIF动图直观展示了扩散模型、MAE可视化等任务的训练效果与生成过程所附说明文档和目录结构也能帮助快速定位代码模块。目前已有186人学习下载适合需要边读边练、快速梳理PyTorch建模流程的开发者。 做深度学习的朋友应该都刷到过龙良曲老师的《PyTorch深度学习实践》系列这套教程最大的特点是代码干净、逻辑清晰每一节讲完理论马上给一段能跑的代码非常适合入门。我自己在带团队带新人的时候也经常把这份代码仓库推给他们当第一份学习材料。不过很多人卡在一个地方视频刷完了、代码也跑通了但真正要拿一个模型出来做自己的任务时又不知道从哪下手。这篇文章就是基于龙良曲的PyTorch学习代码把从环境搭建到Unet图像分割模型复现的完整链路讲清楚包括每一步为什么这么做、中间会踩哪些坑希望能给正在学PyTorch或者准备复现Unet的朋友省点时间。1. 内容整体设计与思路拆解1.1 为什么选择这套代码作为复现基础我第一次看龙良曲的代码仓库时最大的感受是“够用、不绕”。很多教程代码为了展示某个技巧会把简单的网络写得特别复杂动不动就上自定义Layer、钩子函数、分布式封装对新手来说信息量过载。龙良曲的代码不一样线性回归就是线性回归卷积网络就是卷积网络每一段代码只解决一个核心问题变量命名也直白读起来像在跟一个同事对需求。但光有代码还不够学习路径也很关键。我的建议是按“基础模块 → 小型实战 → 语义分割”三步走。基础模块包括张量操作、自动求导、nn.Module的搭建流程这些在龙良曲代码里都有对应章节小型实战可以选择他代码里的MINIST分类或者简单的回归任务把训练循环、loss计算、参数更新的流程跑通最后再上Unet这时候你对PyTorch的API已经有手感了再把Unet拆开看会发现它其实就是“卷积块 池化 上采样 跳跃连接”的组合每一块都能在之前的代码里找到原型。1.2 复现Unet时的核心决策点Unet的复现并不是照抄结构就能出结果的有几个决策点会影响最终效果。第一个是输入尺寸的选择Unet内部做下采样和上采样时特征图尺寸必须对齐所以输入尺寸最好是16的倍数比如256×256或512×512省去对齐的麻烦第二个是损失函数的选择二分类分割常用BCEWithLogitsLoss多分类常用CrossEntropyLoss如果样本类别极度不平衡建议加上Dice Loss做组合损失第三个是数据增强策略随机翻转、随机旋转、亮度抖动这几个操作在医学图像分割里几乎必用能有效提升泛化能力。我当时复现Unet时把龙良曲代码里的训练循环模板直接拿过来改了改发现它的通用性确实好数据加载、模型实例化、优化器配置、epoch循环这套结构换任何模型都适用。所以这篇文章写的是Unet但你掌握的方法论可以迁移到DeepLab、PSPNet甚至Transformer类的分割模型上。2. 环境搭建是第一个坎先把这个搞定2.1 用Anaconda创建隔离环境不管你是Windows、Ubuntu还是CentOS我都建议用Anaconda管理Python环境。很多人在系统Python里直接pip install torch装完发现跟TensorFlow或者其他包冲突最后整个环境乱七八糟代码跑不了还找不到原因。用conda建一个独立环境只需要三条命令conda create -n pytorch python3.10 -y conda activate pytorch pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121这里Python版本选3.10是我实测比较稳的3.8、3.9也能用但新版本PyTorch对Python 3.10的优化更好而且很多第三方库比如albumentations对3.10的支持很完善。CUDA版本的选择要看你本机显卡驱动能支持的最高版本用nvidia-smi查看驱动对应的CUDA版本然后选择不高于它的PyTorch版本。NVIDIA官方文档里写明CUDA版本向下兼容比如驱动支持CUDA 12.2那装cu121、cu118都没问题。2.2 CPU版和GPU版的区别如果只是练手跑跑小数据集CPU版也够用但Unet训练跑起来就慢了一个epoch可能要好几分钟甚至更久。我的建议是只要显卡支持一律上GPU版。判断显卡是否支持CUDA先看是不是NVIDIA显卡然后看驱动是否安装正确。集成显卡和AMD显卡就别想CUDA了直接用CPU版走通流程更重要。离线环境的话比如内网服务器可以先在有网的机器上下载好torch的whl包和依赖包再拷贝进去用pip install --no-index --find-links本地路径安装。很多人问为什么安装pytorch一直报错大多不是版本选错而是网络问题导致下载不完整离线安装能绕开这个坑。2.3 验证环境是否装好装完之后一定要做一件事验证CUDA是否真的可用。很多人只看了torch.__version__就以为装好了结果跑训练时发现一直在用CPU慢得离谱还不知道问题在哪。验证代码很简单import torch print(torch.__version__) # 例如 2.8.0 print(torch.cuda.is_available()) # True 才是真的可用 print(torch.cuda.get_device_name(0)) # 显示显卡型号如果torch.cuda.is_available()返回False通常有三种原因一是PyTorch版本和CUDA版本不匹配重新按对应版本的命令安装即可二是显卡驱动版本太低需要去NVIDIA官网更新驱动三是Windows下没有正确安装NVIDIA驱动核心组件可以重装驱动并勾选“执行清洁安装”。3. 从龙良曲代码里拎出最核心的几个知识点3.1 张量与自动求导PyTorch的“记账本”PyTorch最核心的设计就是用动态计算图记录每一步运算训练时调用backward()梯度就能自动回传到每一个参数上。龙良曲的代码里会专门讲张量和自动求导这块一定要吃透不然后面写自定义网络会非常痛苦。张量操作可以类比成Excel里的单元格引用你在A1写了个公式 B1 C1之后B1变了A1跟着变而且它能记住自己依赖了谁。PyTorch里设置requires_gradTrue相当于告诉框架“这个张量需要被跟踪”之后所有用到它的运算都会被记进计算图。反向传播时框架沿着计算图从最后一层往前推把梯度送到每个参数手里。理解了这个机制你就能明白为什么训练代码里每次都要调用optimizer.zero_grad()——不清零的话上一次batch的梯度会累加到这一次参数更新方向就乱了。3.2 nn.Module的搭建范式和forward的关系龙良曲代码里大量使用nn.Module和nn.Sequential前者是模型基类后者是便捷的容器。很多刚接触PyTorch的人会好奇为什么定义了forward方法调用时直接model(x)就行这是因为nn.Module实现了__call__方法它内部会先执行一些钩子和校验再调用你写的forward。理解这一点很重要如果你在模型内部直接调用self.conv.forward(x)就跳过了nn.Module默认的一些处理比如hook、参数更新检查虽然大多数时候不出错但这是个坏习惯保持使用model(x)的调用风格最稳妥。3.3 训练循环的统一模板龙良曲的代码刷完你会发现所有训练代码都长一个样前向传播、算loss、反向传播、优化器更新。这个模板是通用的复现Unet时直接套用即可。for epoch in range(num_epochs): model.train() for images, masks in dataloader: images, masks images.to(device), masks.to(device) outputs model(images) loss criterion(outputs, masks) optimizer.zero_grad() loss.backward() optimizer.step()这段代码值得反复抄写加理解。model.train()和model.eval()的切换也别忘了它影响的是Dropout和BatchNorm的行为训练模式下这两个层有随机性评估模式下要冻结下来保证结果稳定。4. Unet模型复现全流程4.1 Unet结构拆解编码器、解码器和跳跃连接Unet最初是为了医学图像分割提出的但现在已经泛化到各种像素级预测任务。它的结构像一个U形左边是编码器通过卷积和池化逐渐缩小特征图尺寸、增加通道数提取语义信息右边是解码器通过上采样逐级恢复空间分辨率中间用跳跃连接把编码器的特征图拼接到解码器对应层把小尺度的高频细节传过去弥补下采样丢失的位置信息。没有跳跃连接的Unet退化成一条直筒分割边缘会很糊效果大打折扣。复现时最小组成单元是DoubleConv两次卷积加BatchNorm加ReLU激活。可以用循环生成四个编码器块然后倒序生成对应的解码器块最后用一个1×1卷积把通道数映射到类别数量上。代码如下class DoubleConv(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.conv nn.Sequential( nn.Conv2d(in_ch, out_ch, 3, padding1), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), nn.Conv2d(out_ch, out_ch, 3, padding1), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue) ) def forward(self, x): return self.conv(x)4.2 数据加载是容易忽略的重头戏Unet的输入是图像和对应的掩码标注。PyTorch的Dataset需要实现三个方法init、len__和__getitem。__getitem__返回一对样本和标签。实际写的时候经常碰到图像尺寸不一致的问题这时统一resize到256×256或者512×512即可。数据增强建议在__getitem__里面做这样每个epoch看到的都是“新”的数据相当于免费扩充了训练集。常用的增强包括随机水平翻转、随机旋转、随机缩放、色彩抖动等。图像分割的数据集格式五花八门有的是单通道PNG有的是三通道彩色标注还有的是JSON多边形标注要转换成掩码。我的习惯是先把所有数据统一处理成numpy数组再用Dataset封装避免在训练循环里处理各种格式转换省心很多。4.3 损失函数和评估指标的选择逻辑选了BCEWithLogitsLoss还是CrossEntropyLoss取决于你的任务类别数。二分类分割比如分割前景和背景用BCEWithLogitsLoss就行它内部已经包含了sigmoid操作所以模型输出层别额外加sigmoid多分类分割用CrossEntropyLoss内部包含softmax输出层也别额外处理。有个常见的坑如果训练时loss是降下来了但预测出来的结果全黑或全白大概率是推理时多做了sigmoid或者样本里前景背景比例严重失衡。评估指标方面很多初学者只盯着loss看但loss降不代表分割效果好因为像BCE这种逐像素损失在类别不平衡时会偏向预测多数类。建议加上mIoU均交并比和Dice系数这两个指标。mIoU计算预测和真实标签的交集与并集之比Dice系数在一定程度上等价于F1分数两者在分割任务里是最常用的评估尺度。实测下来我自己的经验是Dice Loss配合CrossEntropyLoss组合使用对小目标的召回率有明显改善。4.4 训练和推理的完整闭环训练时需要把Dataset、DataLoader、模型、优化器、学习率调度器全部串起来。有一个细节容易被忽略DataLoader的num_workers要根据机器情况设置Windows下最好设为0不然容易报多进程相关的错误Linux下可以设成4或8能明显加快数据加载速度。推理阶段和训练有区别。训练时模型处于train模式推理时要切到eval模式且用torch.no_grad()包裹否则每个中间张量都会被记录显存很快爆掉。加上sigmoid处理输出后阈值取0.5得到分割结果再用颜色映射叠加到原图上可视化。整体流程跑通一遍后你会觉得Unet其实没那么高深。5. 常见问题与排查技巧实录5.1 安装和环境类问题速查表以下是几个我遇到最多、也最典型的安装问题整理成表格方便对照。问题现象可能原因解决办法torch.cuda.is_available()返回False驱动版本低或PyTorch的CUDA版本与驱动不匹配更新显卡驱动用nvidia-smi确认支持的CUDA版本重装对应版本的PyTorch安装时提示找不到匹配版本没有指定正确的index-url或Python版本过旧使用官网生成的命令确保Python 3.8以上内网环境无法在线安装无外网权限在能联网的机器上下载whl包及全部依赖离线安装启动器提示PyTorch与驱动版本不符驱动太旧或PyTorch版本过新升级驱动或换低版本PyTorch5.2 训练过程中高频踩坑纪实训练Unet时最常碰到的就是显存不足OOM。解决办法不是一上来就换更大的显卡先做三件事减小batch size、降低输入分辨率、开启混合精度训练。PyTorch原生支持torch.cuda.amp代码改动极小显存占用能降低30%~50%速度还能提升。第二个高频坑是loss变成nan。这个问题90%出在学习率上把学习率从0.001改成0.0005甚至0.0001经常会好很多。剩下的10%可能是数据里有NaN值检查一下输入图像是否有异常像素。第三个坑是模型输出尺寸跟标签对不上。Unet里卷积默认padding0时每经过一次卷积尺寸就会变小解码器恢复时尺寸对不齐最终输出尺寸和原图不一样。解法是卷积层统一加padding1或者输入尺寸设置为16的倍数。5.3 分割效果差的原因排查分割效果不好问题往往不在模型结构而在数据或训练配置上。首先是检查标签是否正确——有没有背景和前景像素标反或者标签和图像没有对齐。其次看损失函数的权重如果前景目标占比太小BCE loss会倾向把所有像素预测为背景此时给正样本加权重或者换Dice Loss效果立竿见影。最后看数据增强是否过于激进旋转角度太大、亮度调整范围过大可能导致模型学不到稳定的特征反而掉点。有一次我复现Unet在自建数据集上跑前两个epoch还不错后面突然loss骤降但验证集指标也在降一看是标签处理写错了位置训练和验证用了不同版本的预处理逻辑。这个问题后来成了团队代码Review的必查项数据前后处理逻辑必须只写一份。写在最后的几点操作体会从龙良曲的入门代码到Unet的成功复现我个人的体会是学习PyTorch不需要一开始就啃源码先把小的细节用熟再把它们拼装成实际的项目框架。复现Unet最有价值的地方在于它把卷积、池化、上采样、跳跃连接、数据增强、损失函数、训练评估这些零散知识点串成了一条完整链路。如果你已经能跑通基础代码下一步可以试着在Unet基础上做改进把编码器换成ResNet或MobileNet做迁移学习或者把跳跃连接改成带注意力的Attention Gate再进一步试试Unet、Unet 3。这些改进方向在GitHub上都有开源实现挑一个数据集从零开始复现一遍收获会非常大。最后分享一个我自己常用的调试技巧训练过程中定期保存验证集里预测效果最差的几张图肉眼观察模型在哪里输错了。这个习惯帮我快速发现过不少问题比只看指标分数直观得多。希望这篇文章能让你在PyTorch学习和Unet复现的路上少走点弯路。本文还有配套的精品资源点击获取
RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进