UNet遥感图像语义分割实战：从结构原理到训练调试全指南

简介基于UNet的遥感图像语义分割毕业设计项目面向计算机相关专业学生用于毕业设计、期末大作业或课程设计均很合适。资源覆盖数据生成、模型搭建、训练与预测推理的完整流程配套的Notebook交互笔记配合train.py、model.py等六个Python脚本便于从数据集构建一直调试到结果可视化。论文部分收录LaTeX源码、图表及参考文献并附郭子睿《基于U-Net网络的遥感图像语义分割》PDF方便对照撰写修改。资源共68个文件含图片、矢量图、字体文件及启动辅助环境的便捷脚本压缩包约47MB。源码已本地编译可运行评审分98分难度适中并通过助教审定已有369人学习适合需要可运行完整项目与论文模板参考的计算机专业学生。1. 为什么UNet做遥感语义分割值得做先搞清楚这个毕设到底在解决什么问题遥感图像语义分割简单说就是给一张卫星或航拍影像里的每个像素打上“道路、建筑、农田、水体、植被”这类地物标签很多人第一反应是UNet已经过时了语义分割不是还有DeepLabV3、SegFormer吗但真要把一个项目从数据、训练、评估到论文完整跑通UNet恰恰是翻车概率最低的选择结构简单、显存友好、训练稳定拿到中等规模的公开遥感数据集上mIoU能到65%—75%这个区间足够支撑一篇合格的毕业设计论文。本文不聊论文里的套话直接讲清楚UNet每个模块在遥感图像上的实际意义以及从原始大图到最终提交的源码和论文素材整条管线怎么搭哪些参数别乱动哪些坑我替你踩过了。适合谁手里有遥感图像源或者用公开数据集、需要快速出结果的你。2. UNet结构为什么适合遥感图像编码器、跳跃连接与三个选型理由2.1 编码器在提取什么遥感影像的层次从像素到地物遥感影像和自然图像最大的不同是地面采样距离GSD也就是一个像素对应地面多少米。0.5米分辨率和10米分辨率的图同一个“道路”在语义上完全不同。UNet的编码器通过四次下采样池化或跨步卷积把输入从512×512一路压到32×32通道数从64涨到512本质上是在堆叠感受野浅层卷积看到的是小范围的纹理和边缘深层卷积看到的是大范围的上下文。对遥感场景来说这个大范围上下文恰好对应“一片农田周边是水体还是建筑”这类判断。下采样次数决定了网络能看到的上下文范围。遥感地物尺度差异极大一棵树的冠幅可能只有几个像素一片农田却能跨几百个像素所以要保证最深层的特征图不至于太小。原始UNet设计里512×512输入对应32×32最深层如果你的输入是256×256建议把下采样次数从4次降到3次否则最深层的特征图只有16×16小目标直接消失。下面这段代码就是UNet里最基础的卷积单元注意padding必须设成1这样特征图尺寸在卷积前后保持不变后面的跳跃连接拼接才不会出维度错乱。import torch import torch.nn as nn class DoubleConv(nn.Module): 两次卷积 批归一化 ReLUUNet 最基础的重复单元 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) class Down(nn.Module): 下采样最大池化 DoubleConv def __init__(self, in_ch, out_ch): super().__init__() self.mpconv nn.Sequential( nn.MaxPool2d(2), DoubleConv(in_ch, out_ch) ) def forward(self, x): return self.mpconv(x)这段代码里有一个新手最容易忽略的点padding1不是随便写的它确保卷积输出尺寸等于输入尺寸这样跳跃连接拼接时通道维度翻倍、空间维度完全对齐。如果你用了padding0特征图每过一层卷积就缩一圈走到解码器拼接时要么报尺寸不匹配的错要么需要强行中心裁剪信息就丢了。另外注意inplaceTrue这个开关它让 ReLU 原地修改输入张量以省显存在深层网络里能省下可观的内存但如果你后续还要用这个张量做别的计算记得把inplace关掉。2.2 跳跃连接为什么是UNet的杀手锏边缘恢复与多尺度特征UNet 最标志性的设计就是编码器和解码器之间的横向跳跃连接。编码器下采样四次之后空间分辨率掉了16倍一个物体的边缘在哪、边界有多锐利这些细节已经在深层特征里被平均掉了。跳跃连接做的事情很简单把编码器第 i 层的特征图直接拼接到解码器第 i 层让分割头同时拿到深层的语义判断和浅层的精细边界。遥感地物的边界往往很清晰道路、屋顶、水体边界锐利边缘模糊一点mIoU 就直接掉两三个点这条连接就是最直接的后悔药。有个经验第一次跑通模型后做一个去掉所有跳跃连接的消融实验你会发现验证集 IoU 下降非常明显特别在道路和建筑物这两个类别上。原因也很直接——这两类物体的形状高度规则边缘信息对正确分类的贡献远大于上下文信息。拼接操作用的是torch.cat通道数翻倍和相加add相比保留了两条独立的信息流解码器可以自己选择侧重哪一路这在遥感这类边界敏感任务里比add效果好。class Up(nn.Module): 上采样 跳跃连接拼接 DoubleConv def __init__(self, in_ch, out_ch): super().__init__() # 转置卷积通道减半尺寸翻倍 self.up nn.ConvTranspose2d(in_ch, in_ch // 2, kernel_size2, stride2) self.conv DoubleConv(in_ch, out_ch) # 拼接后通道 in_ch def forward(self, x1, x2): x1 self.up(x1) # 处理边缘尺寸不一致的常见 bug diff_y x2.size()[2] - x1.size()[2] diff_x x2.size()[3] - x1.size()[3] x1 nn.functional.pad(x1, [diff_x // 2, diff_x - diff_x // 2, diff_y // 2, diff_y - diff_y // 2]) x torch.cat([x2, x1], dim1) return self.conv(x)这里有个很隐蔽的坑转置卷积的输出尺寸不总是精确等于编码器特征图的偶数倍当输入尺寸是奇数时拼接前必须做 padding 对齐否则直接报错。上面这段代码的diff_y、diff_x计算就是干这个的。很多同学在本地跑通、换到自己的数据集上就报 size mismatch九成是这个原因。2.3 选型对比UNet、UNet、DeepLabV3毕设到底选谁很多同学会问一个问题DeepLabV3 用了空洞卷积和 ASPP 模块理论上能捕获多尺度上下文指标上也更高为什么不直接选它我的建议是看你的任务约束毕设要求的是“能稳定复现 能解释清楚 显存不爆”。DeepLabV3 的 ASPP 模块开了多个不同膨胀率的空洞卷积并行显存占用比 UNet 高出一截而且膨胀率、输出 stride 这些超参数对结果影响很大调起来就是玄学。UNet 的归纳偏置更简单强跳跃连接 对称编解码你的论文里每一条设计都有直观解释答辩时也站得住。模型参数量base64显存占用256×256输入遥感小数据集表现训练稳定性调试难度UNet约31M约4GB中上高低UNet约42M约6GB略高中中DeepLabV3约44M约8GB高中高如果你用的是公开遥感数据集数据量通常只有几千张UNet 这类参数量适中的模型反而不容易过拟合。UNet 的密集跳跃连接在这种数据规模下优势有限却多了近一倍的训练时间。如果你觉得 UNet 指标不够亮眼最值得做的改进不是换模型而是换 backbone——把编码器换成预训练的 ResNet34然后接 UNet 的解码器这个操作的收益远大于换 UNet而且代码改动量小后面论文里还能多一个消融实验的对比行。3. 遥感数据准备大图切片、mask转one-hot与增强边界3.1 大图切片为什么不能把整张遥感图直接喂给网络遥感原始影像动辄几千乘几千像素直接输入网络有两个问题一是 GPU 显存放不下二是语义分割模型本质上是滑动窗口式地做局部预测整图输入超出了网络感受野全局信息和局部信息耦合反而让训练不稳定。常见做法是滑动窗口切片把大图切成 256×256 或 512×512 的小块。窗口越小切出的样本越多但上下文越少道路这类长条地物在过小的窗口里会被割裂所以一般取 256×256 起步显存够就上 512×512。import cv2 import numpy as np import os def sliding_crop(image_path, mask_path, out_dir, crop_size256, stride256): 滑动窗口裁剪大图与标签图 stride 小于 crop_size 时输出重叠块重叠能缓解边缘切割问题 img cv2.imread(image_path) # BGR 顺序注意后续别和 RGB 搞混 mask cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) h, w mask.shape index 0 for y in range(0, h - crop_size 1, stride): for x in range(0, w - crop_size 1, stride): img_crop img[y:y crop_size, x:x crop_size] mask_crop mask[y:y crop_size, x:x crop_size] if mask_crop.max() 0 and np.random.random() 0.7: continue # 全背景块按概率丢弃减少无效计算 cv2.imwrite(os.path.join(out_dir, fimg_{index}.png), img_crop) cv2.imwrite(os.path.join(out_dir, fmask_{index}.png), mask_crop) index 1 return index这段代码有个细节值得注意全背景块被以70%的概率丢弃。遥感图的背景比例通常远高于前景如果不做丢弃训练集里90%以上是纯背景块模型很快就学会“全都预测背景”后面再怎么调学习率都救不回来。但也不能全丢否则网络在推理时遇到真实场景里的纯背景区域会乱预测。stride参数默认等于crop_size也就是无重叠切割如果你发现道路在切块边缘频繁断裂可以把stride调成crop_size // 2重叠采样能让道路这类长条目标至少完整落在某一块里。3.2 标签处理灰度掩码转one-hot的正确姿势遥感语义分割的标签通常是灰度 PNG像素值就是类别索引比如0背景、1道路、2建筑、3水体。PyTorch 的CrossEntropyLoss也接受这种类别索引格式不强制转 one-hot。但很多经典评价指标和可视化代码是按 one-hot 设计的而且做数据增强时如果对标签做插值比如随机旋转补边类别索引格式直接插值会产生小数混叠出现 1.4 这种非法类别。这时候就需要先转 one-hot 再做增强增强完用argmax转回索引。def mask_to_onehot(mask_path, num_classes): 灰度掩码 [H, W] 转 one-hot [H, W, num_classes] mask cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) one_hot np.eye(num_classes, dtypenp.uint8)[mask] # 查表法速度远快于循环 return one_hot def onehot_to_mask(one_hot): one-hot [H, W, C] 转回类别索引 [H, W] return np.argmax(one_hot, axis-1).astype(np.uint8)np.eye(num_classes)[mask]这种写法本质是用 mask 的像素值做行索引把每个像素替换成对应的单位向量比 for 循环逐像素赋值快好几个数量级。处理上万张图时这个差异很关键。一个隐藏坑是如果 mask 里出现了num_classes之外的像素值比如标注软件把未标注区域标成255这条查表直接越界报错。所以读取标签后第一件事是检查np.unique(mask)确认类别范围把255统一重映射成对应的忽略类别。3.3 数据增强的边界几何变换随便用色彩变换要克制遥感语义分割里用得最安全的增强是几何类随机旋转90度、水平翻转、垂直翻转这些操作不会改变地物的语义和色彩分布。但遥感影像和自然图像有一个重要差别不同来源的影像色彩差异巨大同一片区域在光学影像和假彩色合成影像下颜色完全不同。做色彩抖动亮度、对比度、饱和度时如果强度过大网络会学到“偏红的就是建筑”这类错误特征换一个数据源直接崩盘。我的习惯是训练时只做几何增强加轻微的尺度抖动推理时完全关闭数据增强。# 训练时用的增强组合只做几何类 import torchvision.transforms as T from PIL import Image train_transform T.Compose([ T.RandomHorizontalFlip(p0.5), T.RandomVerticalFlip(p0.5), T.RandomApply([T.RandomRotation((90, 90))], p0.5), # 只转90度的整数倍 T.ToTensor(), T.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ])注意上面这套 transform 只作用于图像标签必须走另一套只做几何变换的流程且两次变换的随机状态要一致否则图像转了90度而标签没转训练直接报废。实现同步变换可以用同一随机种子或者用torchvision.transforms里的Compose配合自定义函数把图像和标签拼在一起处理。归一化均值方差用 ImageNet 的默认值其实是偷懒做法。遥感影像的像素分布和自然图像差异很大比如近红外通道参与时均值会明显偏移。更稳妥的做法是在自己的数据集上算一遍通道均值和标准差写入代码里这也是论文里可以写进实验设置的一句话成本低收益明确。4. 用PyTorch完整跑通UNet训练模型定义、训练循环与关键参数4.1 从模块到完整模型拼装一个可训练的UNet现在把第2章的模块拼成完整模型。编码器有四层下采样解码器对应四层上采样最后接一个1×1卷积把通道数映射成类别数。基础的base参数设为64也就是第一层卷积的输出通道数后面每一层翻倍。如果你的数据集类别少且图像简单base可以降到32参数量直接缩到原来的四分之一训练速度快很多。class UNet(nn.Module): 完整 UNet四层下采样 四层上采样 跳跃连接 def __init__(self, in_channels3, num_classes6, base64): super().__init__() self.inc DoubleConv(in_channels, base) # 64 self.down1 Down(base, base * 2) # 128 self.down2 Down(base * 2, base * 4) # 256 self.down3 Down(base * 4, base * 8) # 512 self.down4 Down(base * 8, base * 8) # 512最后一层不再翻倍 self.up1 Up(base * 8, base * 4) self.up2 Up(base * 4, base * 2) self.up3 Up(base * 2, base) self.up4 Up(base, base) self.outc nn.Conv2d(base, num_classes, 1) def forward(self, x): x1 self.inc(x) x2 self.down1(x1) x3 self.down2(x2) x4 self.down3(x3) x5 self.down4(x4) x self.up1(x5, x4) x self.up2(x, x3) x self.up3(x, x2) x self.up4(x, x1) return self.outc(x)模型定义里最容易出问题的就是通道数匹配如果你改了base或者增减了层数要一步步数清楚每一层的输入输出通道。一个实用的调试办法是写一段测试代码造一个torch.randn(1, 3, 256, 256)输入打印每一层的输出 shape确认拼接操作没有维度报错。这个检查只要花两分钟能省掉后面训练时反复排查的烦恼。我自己习惯把这段测试代码留在源码的if __name__ __main__里方便以后换数据集时回归验证。4.2 训练主循环损失函数、优化器与学习率配置训练遥感分割模型损失函数的选择要特别谨慎。纯CrossEntropyLoss在不均衡数据集上容易让模型偏向高频类别比如背景占了80%以上模型把所有像素都预测为背景就能拿到很低的 loss。常见做法是交叉熵和 Dice Loss 加权组合Dice Loss 直接优化区域重叠度对不均衡类别更友好。我一般用0.6 * CrossEntropy 0.4 * DiceLoss这个比例在多数遥感数据集上表现稳定。import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader # 损失函数组合缓解类别不均衡 class DiceLoss(nn.Module): def __init__(self, smooth1.0): super().__init__() self.smooth smooth def forward(self, pred, target): pred torch.softmax(pred, dim1) target_onehot nn.functional.one_hot(target, num_classespred.shape[1]) target_onehot target_onehot.permute(0, 3, 1, 2).float() intersection (pred * target_onehot).sum(dim(2, 3)) union pred.sum(dim(2, 3)) target_onehot.sum(dim(2, 3)) dice (2.0 * intersection self.smooth) / (union self.smooth) return 1.0 - dice.mean() criterion nn.CrossEntropyLoss(ignore_index255) dice_loss DiceLoss() def combined_loss(pred, target): return 0.6 * criterion(pred, target) 0.4 * dice_loss(pred, target) # 主训练循环 model UNet(in_channels3, num_classes6, base64).cuda() optimizer optim.AdamW(model.parameters(), lr1e-4, weight_decay1e-4) scheduler optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max100) for epoch in range(100): model.train() total_loss 0.0 for img, mask in train_loader: img, mask img.cuda(), mask.cuda() pred model(img) loss combined_loss(pred, mask) optimizer.zero_grad() loss.backward() optimizer.step() total_loss loss.item() scheduler.step() print(fepoch {epoch} loss {total_loss / len(train_loader):.4f})这段代码里有几个参数直接决定训练成败。学习率1e-4是 UNet 训练比较稳的起点不要一上来就用1e-3很容易在本地数据集上震荡。weight_decay1e-4做 L2 正则防止过拟合。CosineAnnealingLR把学习率从1e-4平滑降到接近0比固定学习率省心很多。如果你换用ReduceLROnPlateau这种按指标衰减的策略记得把监控指标设为验证集 mIoU 而不是 loss——训练 loss 下降不代表验证指标上升。4.3 三个必调参数裁剪尺寸、batch size与类别权重裁剪尺寸直接决定模型的感受野和显存占用。256×256 是绝大多数 GPU 的甜点尺寸RTX 3060 级别跑 batch size 8 没问题512×512 显存占用直接翻四倍batch size 要降到 4 甚至 2。我的经验是如果你的数据集里道路、河流这类长条地物占比高优先上 512×512因为长条目标需要更宽的上下文才能被完整识别用 256×256 会把一条完整道路切成好几段模型很难学出连续性。参数建议值调参方向crop_size256 或 512小目标多就 256长条地物多就 512batch_size8256尺寸/ 4512尺寸显存不够先降 batch再考虑梯度累积base_channels32 或 64数据量小用 32数据量大用 64学习率1e-4loss 震荡就降一半类别权重按频率倒数背景占比过高时启用类别权重是处理不均衡的另一个手段给低频类别更高的 loss 权重等价于强行让模型“重视”少数类。在 PyTorch 里传入CrossEntropyLoss(weighttorch.tensor([...]))即可权重通常设成1 / 类别频率然后归一化。但注意类别权重和 Dice Loss 不要同时拉满两个都强调少数类会把训练过程变得不稳定loss 曲线会像锯齿一样跳动我一般要么用权重交叉熵要么用交叉熵加轻量 Dice二选一。5. UNet遥感分割五大翻车现场现象、原因与排查顺序5.1 翻车现场一模型把整张图预测成背景mIoU 只有个位数现象训练 loss 在下降但验证集输出全黑或全是一个类别mIoU 不到5%。原因类别极端不均衡背景像素占比超过了90%模型发现全预测背景就能把 loss 压得很低梯度更新方向被高频类别主导。这是遥感分割里最高频的翻车原因没有之一。解决先统计数据集的类别像素频率打印每类的占比。然后按 5.3 的方法叠加 Dice Loss或给交叉熵设置类别权重。还有一个非常有效的操作计算每个 batch 里背景和前景的比例在__getitem__里做前景采样权重让每个 batch 至少包含一定比例的前景像素。有人会觉得麻烦但这比调 loss 权重更彻底。5.2 翻车现场二训练 loss 波动剧烈验证集指标上蹿下跳现象loss 曲线不像自然图像训练那样平滑下降而是隔几个 epoch 突然跳高又落回验证集 mIoU 忽高忽低。原因多为两个因素叠加——batch size 太小导致 BN 统计量不稳定遥感图像内部方差又大同一个类别在不同区域色调差异明显梯度方向频繁变化。解决batch size 低于 4 时把BatchNorm2d换成GroupNorm或者干脆冻结 BN 的 running stats。更简单的做法是增大 batch size 或使用梯度累积accumulation_steps4等效于 batch size 翻四倍同时把学习率下调一半。排查顺序建议先看 batch size再看学习率最后检查数据加载顺序有没有打乱——shuffle 被误关也会造成这种震荡。5.3 翻车现场三道路预测断成一截一截边界锯齿明显现象其他类别指标正常唯独道路这类长条状目标的预测结果不连续中间断开边缘像锯齿。原因这类目标的上下文大于裁剪尺寸道路跨越了好几块训练切片模型没有足够的上下文判断断开的区域仍然是道路。另一个原因是转置卷积上采样的棋盘效应直接放大特征图时会留下网格状的伪影。解决裁剪尺寸从 256 提升到 512配合stride重叠采样。同时把转置卷积替换为nn.Upsample(scale_factor2, modebilinear, align_cornersTrue)加普通卷积棋盘效应立刻消失代价是上采样过程丢失一点点细节。对道路这类线状地物还可以在推理后做形态学闭运算先膨胀再腐蚀来填补细小的断裂这个后处理在论文里写一句“应用形态学后处理提升道路连通性”是合法且常见的操作。5.4 翻车现场四训练集指标尚可换一张遥感图效果崩盘现象验证集 mIoU 不错但你从公开数据集里另外下载一幅完全不同来源的遥感图直接推理结果惨不忍睹。原因这是典型的域偏移。不同卫星传感器、不同季节、不同大气条件下的同一地物在影像上的颜色和纹理差异很大。如果训练数据只包含单一来源的影像模型学到的是“这个传感器的色调分布”而不是地物本身的语义特征。解决训练数据尽量混合多个来源的影像哪怕只有几百张也能显著提升泛化。其次是做归一化时用自己的统计数据而非 ImageNet 默认值把每个通道的分布拉平。最后推理阶段做一个简单的直方图匹配让输入图的色彩分布向训练集靠近这个技巧能救回很多分数。5.5 翻车现场五GPU 显存溢出batch size 调到 2 都跑不动现象程序启动后没几分钟就报CUDA out of memory。原因除了输入太大还有一个隐蔽因素——模型在训练模式下会保存中间激活值用于反向传播UNet 的跳跃连接特征图全都被保留在显存里层数越深、通道越多显存占用越大。解决优先检查输入尺寸和 batch size。如果已经很小还爆显存开启torch.utils.checkpoint梯度检查点用部分重计算换显存代价是训练时间增加大约20%。还有一个容易被忽略的点PyTorch 默认会缓存显存分配器某些机器上报 OOM 不一定是真满可以在训练循环里加一句torch.cuda.empty_cache()并在代码开头设置torch.backends.cudnn.benchmark False这个开关在输入尺寸不变时能降低显存碎片化。显存问题的排查顺序输入尺寸 → batch size → 模型通道数 → 梯度检查点按这个顺序走一般都能解决。6. 从模型到毕设交付评估指标、可视化与论文素材组织6.1 三个指标怎么算、怎么写进论文mIoU 是语义分割最核心的指标逐类别算 IoU 再取平均。除了 mIoU遥感领域还常看 Kappa 系数和 F1前者衡量预测与真实标签的一致性程度后者对不均衡类别更敏感。计算这些指标时注意逐类 IoU 里某一类完全没有出现在 GT 中时直接跳过该类别还是记0会影响最终分值写论文时要明确说明你的处理方式。def compute_miou(pred, target, num_classes): 逐类 IoU平均后返回GT 和预测都未出现的类别跳过 ious [] for cls in range(num_classes): p (pred cls) t (target cls) intersection (p t).sum() union (p | t).sum() if union 0: continue ious.append(intersection / union) return np.mean(ious)6.2 结果可视化用叠加图说话论文里放一张纯预测图和一张原图对比审阅人很难直观判断好坏。更有效的做法是生成带透明度的叠加图原图为底、预测结果以半透明色块覆盖并附上 GT 图对比。用matplotlib把三张图并排保存成一张大图每个类别固定一个颜色这样整篇论文所有实验的图表风格统一阅读体验也更好。6.3 把实验数据组织成论文素材一篇合格的分割论文需要三类实验表格各类别 IoU 明细表、消融实验表、与不同模型对比的汇总表。消融实验建议做三组去掉跳跃连接、把 UNet 换成 DeepLabV3、替换 backbone这三组结果就能撑起“分析与讨论”那一节。写实验结果时具体的类别得分远比平均值有说服力一定要在正文里解释为什么某个类别得分低别只放表格等审阅人自己看。把模型预测失败案例汇总成一张典型的 bad case 图并给出你的定性分析这种做法在答辩时是一个非常加分的环节。整套流程跑下来你会发现这个项目最耗时间的不是模型代码而是数据整理和指标调试。我的一个习惯每一组实验跑完立刻把配置文件、指标结果、bad case 截图存在同一个目录下命名带上实验编号和关键参数。到最后写论文时这些素材直接拖进表格就可以用不用回头翻训练日志。做毕业设计最怕的不是不会写代码而是实验做完了数据找不齐。这个习惯帮你省下的时间比任何模型优化技巧都多希望帮到你。本文还有配套的精品资源点击获取