
简介这是一份面向图像语义分割初学者与进阶开发者的实战项目资料围绕UNet结合ASPP模块展开可用于二分类及多类别分割任务适用于医学影像、遥感解译等场景。资源包共10个文件以4个py源码文件为核心涵盖模型定义、数据加载与训练工具等模块另含3个pyc缓存、1个txt依赖清单、1个md说明及1份docx项目说明书压缩包约22KB体积轻便便于快速部署。项目提供从数据预处理、增强、训练到验证评估的全流程方案支持GPU加速、学习率与批次大小等超参数配置并实时记录损失曲线与IoU、Dice等指标自动保存最优权重。读者可据此掌握UNet与ASPP的代码实现与原理讲解理解掩码标签处理、one-hot编码及命令行参数配置等关键细节并在此基础上调整网络深度或引入注意力机制进行性能优化。目前已有115人学习适合希望系统掌握分割实战的开发者参考。1. 从一次遥感图像分割翻车说起unetASPP 到底解决了什么问题去年帮一个做遥感图像语义分割的团队排查模型问题他们用经典 unet 在 512×512 的航拍图上跑建筑物提取交并比卡在 0.71 上不去边缘碎、大建筑内部空洞、相邻建筑粘连成一片。我把他们的解码器特征图拉出来可视化问题很直观下采样 5 次之后最小特征图只有 16×16一个像素对应原图 32×32 的区域那些细长建筑、窄道路在瓶颈层已经被压没了。这不是训练技巧能救的是感受野和空间分辨率的结构性矛盾。unetASPP 这个组合就是冲着这个矛盾去的。unet 的编码器-解码器加跳跃连接负责恢复空间细节ASPPAtrous Spatial Pyramid Pooling空洞空间金字塔池化插在编码器末端用不同空洞率的并行卷积分支在不进一步下采样、不丢分辨率的前提下把感受野撑到多尺度。遥感图像语义分割、医学细胞分割、街景语义分割这类目标尺度跨度大、边界要求高的场景这个组合是性价比很高的改进起点。这篇笔记把项目说明、完整代码、参数设置和踩坑记录一次讲清新手能照着跑通熟手能直接拿去改自己的 backbone。2. unet 与 ASPP 的选型逻辑为什么是这两个拼在一起2.1 unet 的结构瓶颈到底在哪unet 的经典结构是 4 次下采样加 4 次上采样编码器每下采样一次通道翻倍解码器每上采样一次通道减半同层编码器特征通过跳跃连接 concat 到解码器。这套设计在 2015 年提出来时是为医学图像分割服务的当时输入尺寸普遍是 256×256 甚至更小4 次下采样后瓶颈层还有 16×16勉强够用。问题出在输入变大、目标变复杂之后。以 512×512 遥感图为例4 次下采样后瓶颈层 32×325 次下采样后 16×16。一个 16×16 的特征图要表达整张图所有类别的语义每个位置承载的信息量太大小目标直接消失。更麻烦的是unet 瓶颈层只有一层 3×3 卷积堆叠感受野增长是线性的想覆盖整张图的上下文要么加深网络梯度问题、参数量爆炸要么继续下采样分辨率进一步损失。这就是为什么很多团队在 unet 上堆到 5 层下采样后mIoU 反而下降——细节丢得比上下文赚得多。2.2 ASPP 用空洞卷积换感受野的账怎么算ASPP 的核心是空洞卷积Atrous Convolution也叫扩张卷积。普通 3×3 卷积的感受野是 3×3空洞率为 r 的 3×3 卷积感受野变成 (2r1)×(2r1)但参数量不变还是 9 个权重。空洞率 1、6、12、18 的四个并行分支感受野分别是 3、13、25、37再叠加全局平均池化分支就覆盖了从局部纹理到全图上下文的多尺度信息。这里有个容易算错的账空洞卷积的感受野是理论值实际有效感受野受权重分布影响通常只有理论值的 1/3 到 1/2。所以空洞率不是越大越好18 以上的分支在 512×512 输入上贡献很有限反而因为采样点太稀疏引入网格伪影gridding artifact。我一般用 1、6、12、18 这组对应 DeepLabv3 的推荐配置在遥感数据集上实测比 1、3、6、9 高 1.5 个点左右。2.3 把 ASPP 插在 unet 的哪个位置位置选择有三种常见做法插在编码器末端瓶颈层、插在解码器每一层、插在跳跃连接上。我推荐第一种原因有三。第一瓶颈层特征图最小ASPP 的并行分支计算量可控512×512 输入下瓶颈层 16×16四个分支加全局池化的显存开销不到 200MB。第二瓶颈层语义最强ASPP 在这里做多尺度融合相当于给解码器喂了一份带全局上下文的语义先验。第三插在解码器每层会让参数量和显存翻倍收益却递减因为浅层特征本身感受野就小空洞卷积发挥不出优势。具体接法编码器最后一次下采样后特征图通道数设为 512 或 1024先过一个 1×1 卷积把通道压到 256再进 ASPP 的五个分支每个分支输出 256 通道最后 concat 成 1280 通道过一个 1×1 卷积压回 256送入解码器。这个 256 是超参显存紧张就降到 128精度掉 0.5 个点以内。3. 完整代码落地从数据加载到 ASPP 模块实现3.1 环境与依赖版本锁定这套代码在 PyTorch 1.13 CUDA 11.7 上验证过Python 3.9。依赖不多但版本要对齐尤其是 torchvision 和 numpy版本错位会在数据增强那里报奇怪的错。# 建议用 conda 建独立环境避免和系统包冲突 conda create -n unet_aspp python3.9 -y conda activate unet_aspp # 锁定版本torch 和 torchvision 必须配套 pip install torch1.13.1cu117 torchvision0.14.1cu117 -f https://download.pytorch.org/whl/torch_stable.html pip install numpy1.23.5 opencv-python4.7.0.72 pillow9.4.0 tqdm4.65.0逻辑说明torch 1.13.1 对应 torchvision 0.14.1这是官方配套版本混用会出现torchvision::nms找不到的运行时错误。numpy 锁 1.23.5 是因为 1.24 之后移除了np.float别名老代码里如果有np.float会直接崩。opencv 用 4.7 是因为 4.8 改了cv2.imread的默认行为读灰度图会多一个通道。3.2 ASPP 模块的完整实现import torch import torch.nn as nn import torch.nn.functional as F class ASPPConv(nn.Sequential): 单个空洞卷积分支空洞卷积 BN ReLU def __init__(self, in_channels, out_channels, dilation): super().__init__( nn.Conv2d(in_channels, out_channels, 3, paddingdilation, dilationdilation, biasFalse), nn.BatchNorm2d(out_channels), nn.ReLU(inplaceTrue) ) class ASPPPooling(nn.Sequential): 全局平均池化分支把特征压成 1x1 再上采样回去 def __init__(self, in_channels, out_channels): super().__init__( nn.AdaptiveAvgPool2d(1), nn.Conv2d(in_channels, out_channels, 1, biasFalse), nn.BatchNorm2d(out_channels), nn.ReLU(inplaceTrue) ) def forward(self, x): size x.shape[-2:] # 记录原始 H,W用于上采样还原 for mod in self: x mod(x) return F.interpolate(x, sizesize, modebilinear, align_cornersFalse) class ASPP(nn.Module): 空洞空间金字塔池化五个并行分支 def __init__(self, in_channels, out_channels256, rates(1, 6, 12, 18)): super().__init__() # 1x1 卷积分支捕捉局部细节 self.branch0 nn.Sequential( nn.Conv2d(in_channels, out_channels, 1, biasFalse), nn.BatchNorm2d(out_channels), nn.ReLU(inplaceTrue) ) # 三个空洞卷积分支空洞率 6/12/18 self.branch1 ASPPConv(in_channels, out_channels, rates[1]) self.branch2 ASPPConv(in_channels, out_channels, rates[2]) self.branch3 ASPPConv(in_channels, out_channels, rates[3]) # 全局池化分支捕捉全图上下文 self.branch4 ASPPPooling(in_channels, out_channels) # 融合层5 个分支 concat 后压回 out_channels self.project nn.Sequential( nn.Conv2d(out_channels * 5, out_channels, 1, biasFalse), nn.BatchNorm2d(out_channels), nn.ReLU(inplaceTrue), nn.Dropout(0.5) # 防止过拟合小数据集尤其重要 ) def forward(self, x): x0 self.branch0(x) x1 self.branch1(x) x2 self.branch2(x) x3 self.branch3(x) x4 self.branch4(x) out torch.cat([x0, x1, x2, x3, x4], dim1) return self.project(out)逻辑说明ASPPConv里paddingdilation是关键保证空洞卷积输出尺寸和输入一致否则 concat 时尺寸对不上。ASPPPooling的forward里先记录size再池化最后interpolate回原尺寸align_cornersFalse是 PyTorch 官方推荐设为 True 会在边缘产生偏移。project里的Dropout(0.5)在小数据集上能压 2 个点左右的过拟合数据量超过 1 万张可以降到 0.3 或去掉。参数说明in_channels是编码器末端输出通道unet 用 5 层下采样时通常是 512 或 1024。out_channels默认 256显存不够改 128。rates默认 (1,6,12,18)输入尺寸小于 256 时建议改成 (1,3,6,9)否则大空洞率分支采样点会落到特征图外等价于零填充浪费算力。3.3 把 ASPP 嵌进 unet 的完整网络class DoubleConv(nn.Module): unet 的基础卷积块两次 3x3 卷积 BN ReLU def __init__(self, in_ch, out_ch): super().__init__() self.conv nn.Sequential( nn.Conv2d(in_ch, out_ch, 3, padding1, biasFalse), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), nn.Conv2d(out_ch, out_ch, 3, padding1, biasFalse), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue) ) def forward(self, x): return self.conv(x) class Down(nn.Module): 下采样最大池化 DoubleConv def __init__(self, in_ch, out_ch): super().__init__() self.mp nn.MaxPool2d(2) self.conv DoubleConv(in_ch, out_ch) def forward(self, x): return self.conv(self.mp(x)) class Up(nn.Module): 上采样转置卷积 concat 跳跃连接 DoubleConv def __init__(self, in_ch, out_ch): super().__init__() self.up nn.ConvTranspose2d(in_ch, in_ch // 2, 2, stride2) self.conv DoubleConv(in_ch, out_ch) def forward(self, x1, x2): x1 self.up(x1) # 尺寸对齐防止奇数尺寸输入导致 concat 失败 diffY x2.size()[2] - x1.size()[2] diffX x2.size()[3] - x1.size()[3] x1 F.pad(x1, [diffX // 2, diffX - diffX // 2, diffY // 2, diffY - diffY // 2]) return self.conv(torch.cat([x2, x1], dim1)) class UNetASPP(nn.Module): unet ASPP 完整网络4 层下采样 def __init__(self, in_channels3, num_classes2, base_ch64): super().__init__() self.inc DoubleConv(in_channels, base_ch) self.down1 Down(base_ch, base_ch * 2) self.down2 Down(base_ch * 2, base_ch * 4) self.down3 Down(base_ch * 4, base_ch * 8) self.down4 Down(base_ch * 8, base_ch * 16) # 瓶颈层 1024 通道 self.aspp ASPP(base_ch * 16, out_channels256) self.up1 Up(base_ch * 16, base_ch * 8) self.up2 Up(base_ch * 8, base_ch * 4) self.up3 Up(base_ch * 4, base_ch * 2) self.up4 Up(base_ch * 2, base_ch) self.outc nn.Conv2d(base_ch, num_classes, 1) def forward(self, x): x1 self.inc(x) # 1/1 x2 self.down1(x1) # 1/2 x3 self.down2(x2) # 1/4 x4 self.down3(x3) # 1/8 x5 self.down4(x4) # 1/16 x5 self.aspp(x5) # ASPP 在瓶颈层做多尺度融合 x self.up1(x5, x4) x self.up2(x, x3) x self.up3(x, x2) x self.up4(x, x1) return self.outc(x)逻辑说明Up里的F.pad是处理奇数尺寸输入的关键比如输入 513×513下采样 4 次后尺寸不整除上采样回来和跳跃连接对不上不加 pad 会直接报错。base_ch64是标准配置显存够可以提到 96 或 128精度提升有限但显存翻倍。ASPP 的in_channels必须和down4输出通道一致这里是base_ch*161024。参数说明num_classes按数据集类别数设二分类设 2多分类设实际类别数加背景。base_ch是通道基数8GB 显存跑 512×512 输入建议 6416GB 可以上 96。Up里in_ch // 2是转置卷积输出通道保证 concat 后通道数等于in_ch。3.4 训练循环与损失函数选择import torch.optim as optim from torch.utils.data import DataLoader def train_one_epoch(model, loader, optimizer, criterion, device): model.train() total_loss 0 for imgs, masks in loader: imgs, masks imgs.to(device), masks.to(device) optimizer.zero_grad() outputs model(imgs) loss criterion(outputs, masks) loss.backward() optimizer.step() total_loss loss.item() return total_loss / len(loader) # 初始化 device torch.device(cuda if torch.cuda.is_available() else cpu) model UNetASPP(in_channels3, num_classes2, base_ch64).to(device) # 交叉熵 Dice 混合损失类别不均衡时比纯 CE 稳 ce nn.CrossEntropyLoss() def criterion(pred, target): return ce(pred, target) 0.5 * dice_loss(pred, target) optimizer optim.AdamW(model.parameters(), lr1e-3, weight_decay1e-4) scheduler optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max50)逻辑说明AdamW比Adam多了正确的权重衰减实现分割任务上泛化更好。CosineAnnealingLR让学习率从 1e-3 余弦降到接近 0比 StepLR 平滑收敛更稳。混合损失里 Dice 系数 0.5 是经验值正负样本比例超过 1:10 时提到 1.0。参数说明lr1e-3是 AdamW 的常用起点batch size 小于 8 时降到 5e-4。weight_decay1e-4防止过拟合小数据集可以提到 1e-3。T_max50是余弦周期和总 epoch 数一致。4. 避坑与排查unetASPP 训练中最容易翻车的 5 个点4.1 空洞率设太大导致网格伪影现象训练 loss 正常下降但验证集 mIoU 卡在 0.6 左右上不去预测结果出现规律性的网格状条纹尤其是大块同质区域。原因空洞率 18 以上的分支3×3 卷积的 9 个采样点间距太大相邻像素之间没有重叠卷积核实际只采样到特征图的稀疏网格丢失了局部连续性。输入尺寸越小越明显256×256 输入下空洞率 18 的有效感受野已经接近全图但采样密度只有 1/18。解决把rates从 (1,6,12,18) 改成 (1,3,6,9)或者根据输入尺寸动态调整经验公式是最大空洞率不超过输入尺寸 / 16。512×512 输入用 18 没问题256×256 输入降到 12128×128 输入降到 6。4.2 瓶颈层通道数不匹配导致 concat 报错现象RuntimeError: Sizes of tensors must match except in dimension 1报错位置在 ASPP 的torch.cat或解码器的 concat。原因ASPP 的in_channels和编码器末端输出通道不一致。常见于改了base_ch但忘了同步改 ASPP 的入参或者用了预训练 backbone如 ResNet但没对齐通道数。解决在UNetASPP.__init__里把 ASPP 的in_channels写成base_ch * 16和down4的输出绑定。如果用预训练 backbone先打印一次x5.shape确认通道数再传进去。4.3 小数据集上 Dropout 位置放错现象训练集 mIoU 0.95验证集 0.62典型过拟合加了 Dropout 也没改善。原因Dropout 加在了 ASPP 的project之后但解码器的DoubleConv里没有正则化。分割任务的过拟合主要发生在解码器的高分辨率层那里参数量大、特征图大ASPP 在瓶颈层参数量小加不加 Dropout 影响有限。解决在DoubleConv的第二个 ReLU 后加nn.Dropout2d(0.2)解码器每层都加。ASPP 的 Dropout 保留 0.5 不变。数据增强也要跟上随机翻转、旋转、颜色抖动三件套至少上两件。4.4 学习率预热缺失导致早期震荡现象前 5 个 epoch loss 剧烈震荡从 0.8 跳到 2.3 再掉回来mIoU 前 10 个 epoch 几乎不涨。原因ASPP 的五个分支初始化后输出尺度不一致全局池化分支的输出方差比空洞卷积分支大concat 后融合层权重需要时间适应。直接用 1e-3 的学习率早期梯度太大融合层权重被推来推去。解决加 5 个 epoch 的线性预热学习率从 1e-6 线性升到 1e-3再走余弦退火。PyTorch 里用LambdaLR实现或者用torch.optim.lr_scheduler.SequentialLR组合。4.5 验证时忘了切 eval 模式现象验证集 mIoU 比训练集低 10 个点以上但 loss 正常模型看起来没崩。原因model.train()模式下 BatchNorm 用当前 batch 的均值和方差验证集 batch size 通常比训练小统计量偏差大。Dropout 也在 train 模式下随机丢弃验证结果不稳定。解决验证循环开头加model.eval()结尾加model.train()。用torch.no_grad()包住验证循环省显存也提速。这个坑太常见了我见过不止一个团队在这上面浪费一周。5. 进阶技巧用多尺度推理把 mIoU 再提 2 个点训练完之后别急着收工推理阶段还有免费的精度可以捡。多尺度推理Multi-Scale Inference是分割任务里性价比最高的技巧不需要重新训练只在验证和测试时把输入缩放到多个尺度分别预测再取平均。具体做法把原图缩放到 0.75、1.0、1.25、1.5 四个尺度分别送进模型得到四张概率图再缩回原尺寸做平均。代码实现如下def multi_scale_inference(model, img, scales(0.75, 1.0, 1.25, 1.5)): 多尺度推理返回平均后的概率图 model.eval() h, w img.shape[-2:] probs torch.zeros(1, 2, h, w, deviceimg.device) # 2 是类别数 with torch.no_grad(): for s in scales: new_h, new_w int(h * s), int(w * s) # 缩放到目标尺度注意用 bilinear 保持平滑 scaled F.interpolate(img, size(new_h, new_w), modebilinear, align_cornersFalse) out model(scaled) # 缩回原尺寸softmax 后再平均 out F.interpolate(out, size(h, w), modebilinear, align_cornersFalse) probs F.softmax(out, dim1) return probs / len(scales)逻辑说明F.interpolate缩放输入时用bilinear比nearest平滑避免引入锯齿。输出缩回原尺寸后再softmax保证每个尺度的概率分布独立归一化平均才有意义。如果先平均 logits 再 softmax大尺度的过自信预测会主导结果实测比先 softmax 低 0.8 个点。参数说明scales的选择看显存和推理时间预算四个尺度是精度和速度的平衡点。显存紧张可以只留 1.0 和 1.25 两个尺度精度提升减半但推理时间只多一倍。1.5 尺度在 512×512 输入下会变成 768×768显存占用是 1.0 尺度的 2.25 倍8GB 卡上要小心 OOM。还有一个配套技巧是水平翻转推理TTA把原图和水平翻转后的图分别预测翻转回来再平均。和多尺度叠加使用在遥感数据集上实测能把 mIoU 从 0.78 推到 0.81。代价是推理时间变成 8 倍4 尺度 × 2 翻转线上服务要权衡。我自己的习惯是训练阶段只跑单尺度验证选最优 checkpoint最后用多尺度加翻转跑一次测试集出最终指标。这样训练快最终精度也不吃亏。踩过的坑是早期直接在训练循环里加多尺度验证每个 epoch 多花 3 倍时间调参效率极低后来改成每 10 个 epoch 跑一次多尺度才把节奏拉回来。希望帮到你。本文还有配套的精品资源点击获取