简介:面向毕业设计场景的基于UNet遥感图像语义分割完整项目,适合计算机视觉、遥感相关专业学生作为课题参考,尤其适合需要从零搭建分割模型并完成论文写作的读者。资源包共69个文件,约47.26MB,包含Python源码(UNet编码器-解码器架构、数据加载、训练与预测脚本)、Jupyter Notebook交互式实验、遥感图像样本(PNG/BMP)、毕业论文PDF与LaTeX章节源码等,目录结构清晰,并附带README说明,便于按模块检索。已有346人学习下载。内容涵盖数据准备、像素级标注、预处理、UNet网络搭建、反向传播训练以及IoU等指标评估,附带的notebook可逐步演示数据创建、训练与预测过程,帮助理解遥感图像中建筑物、道路等类别分割的实现细节;毕业论文LaTeX源码和PDF提供了完整写作框架,可结合实验数据直接修改使用,并涵盖推理应用阶段的部署思路。对需要快速复现实验、参考完整毕设结构的同学,这份资源兼具代码与文档双重价值。
1. 基于UNet的遥感图像语义分割:这份毕业设计资源到底能帮你什么
如果你正在做“基于UNet的遥感图像语义分割”这个方向的毕业设计,多半需要的不是又一篇理论综述,而是一套能跑通的代码和一条不绕弯路的实操流程。这份UNet-AerialImageSegmentation-master压缩包正好把这条链路补齐了:数据制作、模型训练、预测推理三件套齐全,连毕业论文的LaTeX源码都在里面,适合拿来当baseline再去做改进实验。
遥感图像的语义分割和自然图像不一样:俯瞰视角、目标尺度小、类别极不均衡,原始大图动辄上亿像素,直接整图喂给网络显存必然爆炸。这套资源的核心价值就是解决“遥感图怎么切、怎么标、怎么喂给UNet训练、预测结果怎么评估”这四个具体问题,对应项目里的create_dataset.ipynb、train.ipynb和predict.ipynb,不会让你对着一个空模型发愁。
代码主干是经典UNet,结构清楚,后续换成注意力UNet、ResUNet或加多尺度模块都有操作空间。下面按实际运行顺序,把文件结构、网络细节、训练参数和高频坑位逐个说清楚,希望能帮到你。
2. UNet架构与遥感分割的任务特点:编码器-解码器为什么适合像素级分类
2.1 UNet编码器:下采样与特征提取的逐层递进
UNet最初是为医学图像分割设计的,但很快在遥感语义分割里成为默认baseline。原因在于它的编码器结构恰好匹配遥感目标的多尺度特性:建筑物是一两百像素的规则矩形,道路是细长条,植被是纹理密集的大面积区域,这些目标分别需要在不同感受野下才能被有效识别。编码器做的正是“用空间分辨率换通道数”,每经过一次池化,特征图尺寸减半、通道数翻倍,网络在更抽象的层次上提取语义。
项目里的src/model.py负责组装完整网络,src/cnn.py定义了两个卷积加一次激活的基础卷积块。典型的UNet编码器结构如下:
# src/model.py 中UNet编码器部分的核心逻辑 class UNetEncoder(nn.Module): def __init__(self, in_channels=3, base_channels=64): super().__init__() # 第一层:从RGB输入进入64通道,特征图尺寸不变 self.enc1 = double_conv(in_channels, base_channels) # 第二层开始:每层先下采样,通道数翻倍 self.enc2 = double_conv(base_channels, base_channels * 2) # 第三层:空间尺寸变为输入的1/4 self.enc3 = double_conv(base_channels * 2, base_channels * 4) # 第四层:空间尺寸变为输入的1/8,通道数最大 self.enc4 = double_conv(base_channels * 4, base_channels * 8) def forward(self, x): # 每一层输出都保留,后面解码器的跳跃连接要用 x1 = self.enc1(x) x1_down = F.max_pool2d(x1, kernel_size=2, stride=2) x2 = self.enc2(x1_down) x2_down = F.max_pool2d(x2, kernel_size=2, stride=2) x3 = self.enc3(x2_down) x3_down = F.max_pool2d(x3, kernel_size=2, stride=2) x4 = self.enc4(x3_down) return [x1, x2, x3, x4]这里double_conv是两个3x3卷积的堆叠,卷积后接BatchNorm和ReLU。在原始UNet论文里没有BatchNorm,但遥感图像输入的像素分布波动比医学影像大,加上BN后训练会稳定很多,现在几乎成了标配。
参数上你要关心的只有两个:in_channels和base_channels。in_channels对应输入图像的通道数,普通RGB遥感图设为3,如果是多光谱影像(比如R、G、B、NIR四个波段),就要改成4。base_channels决定整个模型的宽度,基础值64时UNet参数量大约3100万;显存紧张或训练数据少时降到32,参数量会缩到原来的约四分之一,但分割精度也可能掉两三个点。
一个容易被忽略的点是池化方式。这里用的是2x2最大池化,对边缘纹理不敏感,换用步长为2的卷积下采样能保留更多位置信息,代价是计算量上升。我做遥感分割时一般保留最大池化,因为后面还有解码器做精细恢复,池化丢失的细节可以通过跳跃连接补回来。
遥感图的地物尺度差异极大,单靠编码器逐层下采样还是不够。这也是后来很多改进工作会在编码器底部加空洞卷积或特征金字塔的原因,但理解基础UNet的逐层递进逻辑,是看懂那些改进的前提。
2.2 解码器与跳跃连接:恢复空间分辨率的关键设计
编码器把一张512x512的遥感图压缩成64x64的深层特征,如果直接在这上面做分类,输出分辨率完全不够。解码器的工作就是一步步把空间尺寸恢复回去,同时把通道数降下来,最终输出和输入一样大的逐像素分类结果。
解码器的核心操作有两个:转置卷积上采样和跳跃连接拼接。转置卷积把特征图的宽高翻倍,跳跃连接把编码器对应层的输出拼到解码器特征后面,让网络在恢复分辨率时能参考低层的边缘、纹理信息。对遥感分割来说,这一步尤其关键:建筑物边界是否整齐、道路是否连续,很大程度上取决于跳跃连接有没有把底层空间信息传上来。
# 解码器部分:上采样 + 跳跃连接拼接,最后用1x1卷积输出分类 class UNetDecoder(nn.Module): def __init__(self, base_channels=64, num_classes=6): super().__init__() # 转置卷积:通道减半,尺寸翻倍 self.up1 = nn.ConvTranspose2d(base_channels * 8, base_channels * 4, kernel_size=2, stride=2) # 拼接后通道数 = 编码器enc3输出 + up1输出 self.dec1 = double_conv(base_channels * 8, base_channels * 4) self.up2 = nn.ConvTranspose2d(base_channels * 4, base_channels * 2, kernel_size=2, stride=2) self.dec2 = double_conv(base_channels * 4, base_channels * 2) self.up3 = nn.ConvTranspose2d(base_channels * 2, base_channels, kernel_size=2, stride=2) self.dec3 = double_conv(base_channels * 2, base_channels) # 最后用1x1卷积把通道压缩到类别数 self.final_conv = nn.Conv2d(base_channels, num_classes, kernel_size=1) def forward(self, enc_features): x1, x2, x3, x4 = enc_features # 最深层先上采样,再和对应编码层特征拼接 x = self.up1(x4) x = torch.cat([x3, x], dim=1) x = self.dec1(x) x = self.up2(x) x = torch.cat([x2, x], dim=1) x = self.dec2(x) x = self.up3(x) x = torch.cat([x1, x], dim=1) x = self.dec3(x) return self.final_conv(x)拼接的维度要算清楚:torch.cat([x3, x], dim=1)是把通道维度拼接,所以要求两路特征图的空间尺寸完全一致。这也是为什么我把每层下采样都设置在尺寸减半后立即进行编码,只要池化步长和卷积填充保持一致,解码器对应层就能对上。
上采样部分的实现选择是个容易翻车的细节。ConvTranspose2d如果kernel_size设置不当,会产生棋盘格伪影,在遥感图里表现为细碎的方块纹理,对道路、房屋边界的预测干扰很明显。常见的替代方案是先用F.interpolate(scale_factor=2, mode='bilinear')做双线性上采样再接普通卷积,效果更平滑。我看src/model.py里用的还是转置卷积,这是因为原始UNet论文就是这个写法。如果你发现预测结果有明显的格子纹理,优先把转置卷积换成“上采样+卷积”的组合,棋盘伪影基本能消除。
另外注意num_classes不只决定最后一层输出通道,还会影响损失函数的计算。如果数据集是6类(背景、建筑、道路、植被、水体、其他),这里就是6;类别数设错,训练时标签索引超出范围会直接报错。
2.3 遥感图和自然图分割的差异:任务特性决定了UNet的网络设计
在VOC或COCO上做语义分割,目标通常是画面中央的物体,占比大、边界清楚。遥感图像完全不是这样:地物是从高空俯视,目标密集且相互遮挡,建筑物屋顶和道路颜色可能相近,树冠阴影会让边界变得不明确。更麻烦的是类别分布极度不均,一大块遥感图里可能95%都是背景或植被,建筑物只占几个百分点。
这意味着你不能直接把自然图像语义分割的代码拿过来就用。数据预处理上,自然图像分割通常整图resize到固定尺寸,遥感图因为原图太大,必须切成patch训练;类别不平衡上,普通交叉熵会让模型学会“全部预测为背景”这种作弊策略,必须加类别权重或改用Dice Loss;评估上也一样,只看整体像素准确率没有意义,因为全部预测背景也有95%的准确率,要按类别分别计算IoU才看得出模型是否真的学到了地物结构。
遥感语义分割和遥感图像目标检测也是两回事。目标检测输出的是边界框,告诉你“哪里有个房子”;语义分割输出的是逐像素类别标签,告诉你“每一个像素属于哪一类”。后者对边界精度要求高得多,这也是为什么UNet这类编码器-解码器结构比单纯的分类网络更适合这个任务——它在恢复分辨率的过程中保留了边界细节。
理解了这些差异,你再看后面数据集制作和训练参数的选择,就会明白每一步都是为了应对遥感图像的特殊性,而不是无脑套用通用流程。
3. 从压缩包到可训练数据:项目结构、环境准备与标注处理
3.1 解压后先看什么:文件结构与启动脚本
拿到压缩包先别急着跑训练,把文件结构扫一遍,弄清楚每个文件的职责。解压后核心内容如下:
UNet-AerialImageSegmentation-master/ ├── README.md # 项目说明和运行指南 ├── start_jupyter.ps1 # 一键启动Jupyter的PowerShell脚本 ├── start_tensorboard.ps1 # 一键启动TensorBoard的脚本 ├── src/ │ ├── __init__.py │ ├── cnn.py # 基础卷积块:double_conv等 │ ├── model.py # 完整UNet模型定义 │ ├── data.py # Dataset类、数据加载、预处理 │ ├── train.py # 训练主逻辑 │ └── utils.py # 评估指标、可视化等工具函数 ├── create_dataset.ipynb # 数据集制作:裁剪、标注映射、划分 ├── train.ipynb # 训练流程的交互式版本 ├── predict.ipynb # 推理和预测结果可视化 ├── demo/ # 示例图片或演示输出 ├── 基于U-Net网络的遥感图像语义分割_郭子睿.pdf # 论文PDF └── 毕业论文/ # LaTeX全文:chap1.tex ~ chap5.tex两个PowerShell脚本值得先看。start_jupyter.ps1通常做两件事:激活项目依赖的conda环境,然后启动Jupyter Notebook服务。start_tensorboard.ps1则是启动TensorBoard用于监控训练曲线,脚本里一般会指定日志目录和端口。我在Windows上第一次运行这类脚本时遇到过执行策略拦截,PowerShell默认不允许运行未签名的.ps1文件,解决办法是在当前终端窗口临时放开限制:
# 只对当前PowerShell窗口生效,不改系统设置 Set-ExecutionPolicy -Scope Process Bypass cd UNet-AerialImageSegmentation-master .\start_jupyter.ps1README.md里通常会写依赖的环境版本,我建议先把PyTorch、torchvision、numpy、opencv-python装齐。CUDA版本和PyTorch版本不匹配是最常见的环境问题,推荐用conda创建独立环境,不要直接装在base环境里。我一般这样做:
conda create -n unet python=3.8 conda activate unet conda install pytorch==1.12.1 torchvision==0.13.1 cudatoolkit=11.3 -c pytorch pip install opencv-python pillow matplotlib jupyter tensorboard注意PyTorch版本和CUDA版本要对应,不能随便装最新的PyTorch然后配旧的CUDA驱动。检查CUDA是否可用,在Python里执行torch.cuda.is_available(),返回True再往下走。
查看论文的LaTeX目录还有个额外好处:chap1.tex到chap5.tex是完整的毕业论文正文,Figures目录里是插图。如果你需要写“研究现状”或“实验结果分析”章节,这些tex文件可以直接参考格式,省去从零搭LaTeX框架的时间。
3.2 创建自己的遥感分割数据集:create_dataset.ipynb实操
语义分割数据集制作是整个项目里最费时间的一步,也是最容易出错的一步。create_dataset.ipynb的作用就是把原始遥感大图和标注文件转化成适合UNet训练的patch集合。
先明确标注文件的格式。遥感分割的标签通常是单通道PNG,每个像素的灰度值代表类别索引:0是背景,1是建筑物,2是道路,3是植被,以此类推。也有用RGB彩色编码的标注(比如红色代表建筑物、绿色代表植被),这种情况需要先做颜色映射,把三通道的掩码图转成单通道索引图:
# 将RGB标签图转换为单通道类别索引图 def rgb_to_classmask(rgb_label, color_map): # color_map: {0: (0,0,0), 1: (255,0,0), 2: (0,255,0)} h, w = rgb_label.shape[:2] class_mask = np.zeros((h, w), dtype=np.int64) for cls, color in color_map.items(): # 逐类别匹配颜色,生成布尔掩码后写入类别索引 match = np.all(rgb_label == color, axis=-1) class_mask[match] = cls return class_mask逻辑说明:颜色映射要保证每个类别的RGB值互不相同,且文件里不能出现映射表之外的“脏颜色”。我见过一个数据集在标注时混入了抗锯齿边缘,导致部分像素的颜色在两个类别之间,映射后变成随机类别。解决办法是在生成掩码前先用cv2.medianBlur做一次轻微滤波,或者直接对颜色做最近邻匹配,只保留严格的预设颜色。
接下来是裁剪。遥感大图可能达到5000x5000甚至更大,直接整图进网络不可能,要滑窗裁剪成统一大小的patch。项目里create_dataset.ipynb的核心就是这一步:
# 使用滑动窗口把大图和对应标签裁剪成训练patch def crop_to_patches(image_path, label_path, patch_size=256, stride=128, out_dir='train'): img = cv2.imread(image_path) lbl = cv2.imread(label_path, cv2.IMREAD_UNCHANGED) # 单通道,不能丢通道信息 h, w = img.shape[:2] idx = 0 for y in range(0, h - patch_size + 1, stride): for x in range(0, w - patch_size + 1, stride): img_patch = img[y:y+patch_size, x:x+patch_size] lbl_patch = lbl[y:y+patch_size, x:x+patch_size] # 过滤掉标签全是背景的patch,避免训练样本极度偏向背景 if (lbl_patch == 0).mean() > 0.95: continue cv2.imwrite(f'{out_dir}/img_{idx:05d}.png', img_patch) cv2.imwrite(f'{out_dir}/lbl_{idx:05d}.png', lbl_patch) idx += 1参数说明:patch_size=256是平衡显存和感受野的常用值,512能看到更大的上下文但显存占用直接翻几倍;stride=128是patch大小的一半,相邻patch有一半重叠,相当于做了数据增强,还能避免推理时拼接处的边界不连续问题。如果你后面准备改大模型或加注意力模块,建议训练patch控制在256,不然4GB显存会很吃力。
过滤全背景patch这一步非常关键。遥感图像里背景占比极高,如果不过滤,训练集里可能会有超过一半的patch没有前景目标,模型会被大量无意义的背景样本淹没。但过滤比例也不要太激进,控制在95%以上才丢弃,否则会把少量存在很小目标的patch也丢掉。
裁剪完成后的数据划分也要注意:训练集、验证集、测试集必须从不同的原始大图里划分,不能在同一张大图的不同patch里随机分。因为相邻patch有重叠,如果同源的patch既进了训练集又进了测试集,指标的虚高会让你的论文数据被质疑。我一般先给原始大图编号,按编号分“训练大图、验证大图、测试大图”,再分别裁剪。
3.3 数据预处理与增强的参数选择
数据集制作完成后,训练前还有一道预处理流程:归一化、通道顺序调整、数据增强。这一步写在src/data.py的Dataset类里。
归一化不能想当然。常见错误是直接把像素除以255,但如果原始图是16bit的TIFF(遥感图像很常见),最大像素值是65535,除以255的结果完全不正确。正确做法是先探测图像位深,16bit图先转为8bit再进入归一化流程:
# 统一转8bit并归一化到[0,1],再按数据集的均值和标准差做标准化 def normalize(image): if image.max() > 255: image = (image / 65535 * 255).astype(np.uint8) # 16bit转8bit image = image.astype(np.float32) / 255.0 mean = np.array([0.485, 0.456, 0.406]) std = np.array([0.229, 0.224, 0.225]) image = (image - mean) / std return torch.from_numpy(image).permute(2, 0, 1)注意这里顺手反转了通道维度,从[H, W, C]转到PyTorch需要的[C, H, W]。这个细节经常被忽略,导致训练时报维度错误。
数据增强方面,遥感图有一个天然优势:没有严格的“上下方向”概念,旋转和翻转不会破坏语义。这个特性让数据增强手段比自然图像丰富得多。推荐使用水平翻转、垂直翻转和90度旋转的组合:
# Dataset类里的增强逻辑:只做几何变换,不改变像素值 if self.augment: if random.random() < 0.5: img = np.flip(img, axis=1) # 水平翻转 lbl = np.flip(lbl, axis=1) if random.random() < 0.5: img = np.flip(img, axis=0) # 垂直翻转 lbl = np.flip(lbl, axis=0) k = random.choice([0, 1, 2, 3]) if k > 0: img = np.rot90(img, k) # 旋转90度的整数倍 lbl = np.rot90(lbl, k)几何增强不做插值,所以不会引入标签错位的问题。不推荐在训练时做随机缩放,因为语义分割要求像素级对齐,缩放过程中图像和标签的插值方式不一致就会导致边界偏移。
提示:增强代码里图像和标签的所有变换必须相同,不能只翻转图像不翻转标签,否则模型会在错误的监督信号下训练。这一点我在初学时踩过,loss显示正常,但可视化预测时分割结果和原图完全对不上。
预处理完成后,建议从数据集里抽样几个patch,把标签重叠显示到图像上检查一遍。打开create_dataset.ipynb最后的效果展示部分,能看到mask和原图的叠加结果,这一步能提前发现标注类别错乱、颜色映射错误、裁剪错位等问题,比训练完再看结果省时得多。
4. 训练与推理全流程:train.ipynb和predict.ipynb逐段拆解
4.1 网络模型定义与初始化细节
数据准备好后,进入训练阶段。train.ipynb的第一步是实例化模型。项目里的src/model.py提供了一个完整的UNet类,它把编码器、解码器组装到一起,对外只暴露in_channels、num_classes、base_channels三个参数:
from src.model import UNet model = UNet(in_channels=3, num_classes=6, base_channels=64) model = model.cuda()模型初始化方式对训练稳定性有影响。PyTorch默认的初始化在深层网络上偶尔会带来前期loss不稳定的问题。如果UNet里的卷积层使用了Kaiming初始化,训练会更快进入稳定状态。我通常在构建完模型后挂一层递归初始化:
def init_weights(m): if isinstance(m, nn.Conv2d): nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') if m.bias is not None: nn.init.zeros_(m.bias) model.apply(init_weights)这里fan_out模式适合ReLU激活的卷积层,能保证前向传播的方差稳定。BatchNorm层的权重默认是1、偏置是0,一般不用额外处理。如果之前用别人训练好的权重做微调,就不要执行这步初始化,否则会把加载的权重覆盖掉。
类别数要和数据集的标签索引对齐。比如你最终决定只需要两种类别(建筑物和背景),num_classes就要改成2,且数据集的标签里不能出现2以上的索引值。训练时如果遇到“Target 3 is out of bounds”这类错误,说明标签里有超出类别总数的像素值,需要回到数据处理那一步做修正。
4.2 训练循环、损失函数与优化器配置
训练主逻辑在src/train.py里,train.ipynb是它的交互式版本。核心训练循环不长,但有几个参数值得认真调。
损失函数的选择直接关系到遥感分割的成败。直接使用普通交叉熵,在类别不均衡的数据集上模型很快就会退化到“全背景”策略。项目里提供了加权交叉熵的选项,权重根据训练集类别频率计算:
# 统计训练集所有标签的类别分布 hist = np.bincount(all_labels.flatten(), minlength=num_classes).astype(np.float32) freq = hist / hist.sum() # 中位数频率平衡:少数类权重高,多数类权重低 median = np.median(freq[freq > 0]) class_weight = median / (freq + 1e-8) criterion = nn.CrossEntropyLoss(weight=torch.from_numpy(class_weight).float().cuda()) optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-4)逻辑说明:class_weight是每个类别的损失权重。背景像素占比高,它的权重就被压得很小;建筑物占比低,权重就放大。这样模型即使预测错了背景也不会被无意义地惩罚,每次反向传播时少数类能拿到更可观的梯度。1e-8是防止某些类别频次为0导致除零错误。
优化器我一般选Adam起步,学习率1e-3是比较安全的起点。UNet这类全卷积网络如果用SGD+momentum(lr=0.01, momentum=0.9)也能收敛,但学习率要手动调整epoch间衰减策略,对初学者不太友好。Adam配合ReduceLROnPlateau就能自动降学习率:
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( optimizer, mode='max', factor=0.5, patience=5, verbose=True )mode='max'表示监控的指标越高越好,这里监控验证集的mean IoU;patience=5表示连续5个epoch没提升就降低学习率。训练循环如下:
for epoch in range(epochs): model.train() total_loss = 0.0 for images, labels in train_loader: images = images.cuda() labels = labels.cuda() # 注意labels形状是[B,H,W],不是独热编码 outputs = model(images) # 输出形状[B,num_classes,H,W] loss = criterion(outputs, labels) optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=5.0) optimizer.step() total_loss += loss.item() # 每个epoch结束做一次验证 val_iou = validate(model, val_loader) scheduler.step(val_iou)两个关键细节:一是labels保持[B, H, W]的形状,不要画蛇添足做one-hot编码,CrossEntropyLoss内部会处理;二是梯度裁剪clip_grad_norm_,max_norm设为5.0,能防止个别batch的异常样本导致loss暴涨。我在训练初期遇到过loss突然跳到几百的情况,加上梯度裁剪后再没出现过。
Batch size需要根据显存调整。项目里默认的batch size是8,在6GB显存上跑512x512输入会OOM,如果报CUDA out of memory,先把patch_size降到256,batch size降为4。训练样本少的时候batch size太小,BatchNorm的统计量会不稳定,这是另一个坑,后面专门说。
4.3 评估指标计算:IoU怎么算才有意义
遥感语义分割不能用“准确率”作为核心指标。想象一张图里背景占97%,模型把所有像素都预测成背景,准确率是97%,但实际一个建筑物都没分割出来。所以论文里必须报IoU(Intersection over Union)和mean IoU。
IoU的计算逻辑并不复杂,但实现时有一些容易出错的边界情况:
# 逐类别计算IoU,最后返回mean IoU def compute_iou(pred, label, num_classes): ious = [] for cls in range(num_classes): pred_mask = (pred == cls) # 预测为该类的像素 label_mask = (label == cls) # 真实属于该类的像素 intersection = (pred_mask & label_mask).sum() union = (pred_mask | label_mask).sum() if union == 0: # 该类别在样本中不存在,跳过不参与平均 continue ious.append(intersection.item() / union.item()) return sum(ious) / len(ious) if ious else 0.0逻辑说明:这里的pred是模型输出经argmax(dim=1)后得到的类别索引图,label是真实标注图。每个类别的IoU等于“预测为该类且真实为该类的像素数”除以“两者合计占用的像素数”。这个指标对边界敏感,预测边界偏几个像素都会让IoU明显下降。
计算时的边界陷阱是union == 0的情况。如果某个类在验证集的当前样本中完全不出现,直接计算会得到除零错误或把它算成0分,从而拖低mean IoU。正确做法是跳过不存在的类别,只对实际出现的类别求平均。如果所有类别都不存在(这种情况只会在空图上出现),函数返回0.0而非程序崩溃。
评估时还要注意模型处于评估模式:
def validate(model, val_loader): model.eval() iou_list = [] with torch.no_grad(): for images, labels in val_loader: pred = model(images.cuda()).argmax(dim=1).cpu().numpy() iou_list.append(compute_iou(pred, labels.numpy(), num_classes)) model.train() return np.mean(iou_list)model.eval()会把BatchNorm和Dropout切换到推理模式。如果漏掉这步,BatchNorm会继续用batch统计量而不是全局统计量,验证结果会比实际偏低,而且时高时低不稳定。
除了mean IoU,建议把每个类别的IoU单独打印出来。mean IoU相同的情况下,建筑IoU和道路IoU可能差异很大。论文里如果能分析“模型在道路类上预测结果差是因为道路细长、边缘占比高”,比只报一个平均值要有说服力得多。
4.4 用predict.ipynb做推理与结果可视化
训练完成后,predict.ipynb负责加载权重、对测试图做推理、把预测结果可视化保存。这里有个遥感图像特有的处理:推理时输入不是单个patch,而是一整张大图,必须用滑动窗口方式逐块预测再拼接。
拼接策略直接影响最终分割图的质量。如果滑窗步长等于patch大小,相邻patch没有重叠,模型在不同patch上的预测差异会在拼接处形成明显的接缝。解决方法是采用重叠滑窗,只取每个patch中心区域的预测结果填充到大图上:
# 滑窗推理:stride设为patch_size的一半,重叠区域只取中心部分 def predict_full_image(model, full_img, patch_size=256, stride=128): h, w = full_img.shape[:2] pred_map = np.zeros((h, w), dtype=np.uint8) for y in range(0, h, stride): for x in range(0, w, stride): patch = full_img[y:y+patch_size, x:x+patch_size] # 边界处不足patch尺寸时,复制边缘补齐 if patch.shape[0] < patch_size or patch.shape[1] < patch_size: tmp = np.zeros((patch_size, patch_size, 3), dtype=np.uint8) tmp[:patch.shape[0], :patch.shape[1]] = patch # 复制边缘行/列 if patch.shape[0] < patch_size: tmp[patch.shape[0]:] = patch[-1] if patch.shape[1] < patch_size: tmp[:, patch.shape[1]:] = tmp[:, patch.shape[1]-1:patch.shape[1]] patch = tmp input_tensor = normalize(patch).unsqueeze(0).cuda() with torch.no_grad(): logits = model(input_tensor) pred = logits.argmax(dim=1).squeeze(0).cpu().numpy() # 只取当前滑窗中心的有效区域,避免重叠区域多次覆盖 y_end = min(y + patch_size, h) x_end = min(x + patch_size, w) pred_map[y:y_end, x:x_end] = pred[:y_end - y, :x_end - x] return pred_map逻辑说明:重叠滑窗会让同一像素被预测多次,代码采用后写覆盖的方式,等效于只信任最后一块patch的预测,比直接平均更简单。如果要更精细,可以针对重叠区做加权平均,中心权重高、边缘权重低,但代码复杂度会高一些。大多数毕业设计用后写覆盖就足够了。
推理完成后要把类别索引的可视化。彩色mask叠加到原图上,保存成对比图作为毕业论文的“实验结果”部分:
# 将单通道的预测索引转成彩色图,再与原始图混合显示 color_map = { 0: [0, 0, 0], # 背景 1: [128, 0, 0], # 建筑 2: [0, 128, 0], # 植被 3: [0, 0, 128], # 道路 4: [128, 128, 0], # 水体 5: [128, 0, 128], # 其他 } pred_color = np.zeros((pred_map.shape[0], pred_map.shape[1], 3), dtype=np.uint8) for cls, color in color_map.items(): pred_color[pred_map == cls] = color overlay = cv2.addWeighted(full_img, 0.6, pred_color, 0.4, 0)addWeighted是OpenCV的图像混合函数,参数0.6和0.4控制原图和预测图的透明比例。建议原图权重稍高,这样能看清预测边界和真实地物是否贴合。输出图像时注意,cv2.imwrite保存的是BGR顺序,如果之前在数据预处理里转过RGB顺序,要再转回来,否则可视化出来的颜色会整体偏蓝偏橙。
5. UNet训练遥感分割的常见问题与避坑指南:五个高频踩坑记录
以下五条问题是我在跑这套遥感UNet代码时实际遇到过、并且反复被朋友问到的,每一条都按“现象、原因、解决”记录,方便你对照排查。
5.1 训练loss不降反升的排查思路
现象是训练前几轮loss在2.5左右波动,后续不是下降而是缓慢爬升,验证集的IoU始终为0。
先检查标签里有没有非法像素值。遥感标注图有时会带边界线或未标注区域,像素值可能是255或类别数之外的数字,CrossEntropyLoss计算时遇到超出范围的target会报错或产生NaN梯度。我排查过很多次,最隐蔽的情况是标签文件中混入了255的“忽略区”,让loss在正常类别和无效类别之间反复横跳。
然后是归一化统计量的问题。遥感图像的光谱分布和自然图像差异很大,直接用ImageNet的mean和std有可能让输入分布偏移,导致训练初段loss偏高。建议在src/data.py里加一个统计函数,从训练集抽样200张图计算自己的mean和std,再传入归一化。
评估思路按顺序走:先打印标签像素值的分布,确认最大索引小于num_classes;再打印输入图像的均值方差,确认归一化后数据分布不是极端偏移;最后把学习率降到1e-4跑20步,看loss是否能下降。如果学习率降到1e-4依然不降,问题大概率在数据或标签而不在优化器。
5.2 显存溢出与batch size、图像尺寸的取舍
第一次跑train.ipynb,最容易报的就是CUDA out of memory。UNet虽然是经典结构,但对显存并不友好,每一层编码器和解码器的特征图都要保留在显存里用于反向传播,512x512输入、batch size 8、base_channels 64,8GB显存会被打满并直接OOM。
先明确一个原则:优先降patch_size而不是batch size。patch_size从512降到256,分辨率变成四分之一,显存占用大幅下降;batch size降至2也能缓解显存,但过小的batch会趋势BatchNorm统计不稳定。显存还是不够时,用梯度累积来模拟更大的batch:
# 显存只够batch size=2,但想要等效batch=8的效果 accumulation_steps = 4 batch_size = 2 for i, (images, labels) in enumerate(train_loader): images, labels = images.cuda(), labels.cuda() outputs = model(images) # 除以accumulation_steps是为了累积时的loss量级和正常batch一致 loss = criterion(outputs, labels) / accumulation_steps loss.backward() if (i + 1) % accumulation_steps == 0: optimizer.step() optimizer.zero_grad()参数说明:实际batch size是2,每4次反向传播才更新一次参数,等效于batch size 8的梯度累计。但BatchNorm的统计仍然基于每次前向的batch size 2,效果和真正的batch 8有细微差别。如果任务对精度要求高,建议改成InstanceNorm或GroupNorm来替代BatchNorm,避免小batch下的统计振荡。
显存优化的次序建议:先降patch_size到256,再减batch size到4,再开torch.backends.cudnn.benchmark=True加速卷积,最后才考虑梯度累积。修改完显存还不够,就缩小base_channels到32,这是对模型结构的影响最大但降低显存也比较显著的手段。
5.3 标注掩码与图像尺寸不对齐导致的分割边界错位
现象是训练正常、IoU正常,但可视化预测结果时发现分割边界整体偏移了几个像素,或者建筑边缘出现了双层叠影。
根本原因多半是数据生成阶段图像和标签的处理方式不一致。典型的错误是在裁剪或缩放时用了不同的插值方法:图像用cv2.resize默认的双线性插值,标签也用了双线性插值,于是标签的边界被模糊、类别索引被插值污染,产生了0到1之间的小数,再经过round或取整操作,边界错位就产生了。
标签数据只能使用最近邻插值:
# 缩放标签时强制使用INTER_NEAREST lbl_resized = cv2.resize(lbl, (patch_size, patch_size), interpolation=cv2.INTER_NEAREST) # 图像可以使用双线性或三次插值 img_resized = cv2.resize(img, (patch_size, patch_size), interpolation=cv2.INTER_LINEAR)还有一个隐蔽的坑:读取标签时用了默认的cv2.imread(path),此时三通道图像会被读成BGR三通道,灰度值相同的mask被复制成了三个通道,之后在计算类别时会产生错误。读取单通道标签图必须加cv2.IMREAD_UNCHANGED或cv2.IMREAD_GRAYSCALE。
判断是否存在该问题的快速方法是随机选3到5个patch,把标签以半透明方式叠到原图上,观察建筑的边缘是否贴合屋顶轮廓。只要有一个patch出现错位,就要回头检查数据生成的完整链路,不要尝试在训练阶段弥补。
5.4 类别不平衡让模型只预测背景怎么办
现象是从第一个epoch开始,loss就在下降,但预测结果整片都是背景色,mean IoU一直是0,个别类别的IoU完全没有输出。
原因是遥感数据集中前景类别占比太低。一个典型的城市遥感图像中,建筑物和道路可能只占总像素的10%以下,普通交叉熵在训练时被大量背景像素主导,模型发现“全部预测背景”已经是较低loss的选择,于是陷入局部最优。
加权交叉熵是最快见效的手段,src/data.py里可以按中位数频率平衡法计算类别权重:
# 计算类别权重:中位数频率平衡 hist = np.bincount(all_labels.flatten(), minlength=num_classes).astype(np.float32) freq = hist / hist.sum() median = np.median(freq[freq > 0]) class_weight = median / (freq + 1e-8) # 少样本类得到大权重另一个有效手段是混合Dice Loss。Dice系数衡量的是区域重叠率,对像素数量不敏感,两个类别的像素占比相差30倍,Dice系数会促使模型同时优化它们:
def dice_loss(pred, target, num_classes): # pred是softmax后的概率[B,C,H,W],target是类别索引[B,H,W] loss = 0.0 for cls in range(num_classes): p = pred[:, cls, :, :] t = (target == cls).float() intersection = (p * t).sum() dice = (2.0 * intersection + 1.0) / (p.sum() + t.sum() + 1.0) loss += 1.0 - dice return loss / num_classes # 混合损失:交叉熵 + 0.3 * Dice Loss total_loss = ce_loss + 0.3 * dice_loss(outputs, labels, num_classes)参数说明:Dice Loss前加的0.3是权重系数,太大(比如1.0)会让训练早期梯度变化过快,太小(0.1以下)又起不到缓解类别不平衡的作用。这两个方法配合使用,在多数遥感数据集上能让前景类别的IoU提升5到10个点。
5.5 模型保存与恢复:避免训练到一半白费
训练到20个epoch,mean IoU已经从0.2涨到0.6,笔记本却因为显存不足导致Jupyter内核崩溃,重启后所有训练进度归零,这是最让人崩溃的场景。
只保存最后的权重文件远远不够。训练状态包含模型权重、优化器状态、学习率调度器状态和当前epoch,Pytorch提供了一条命令保存全部信息:
torch.save({ 'epoch': epoch, 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'scheduler_state_dict': scheduler.state_dict(), 'best_iou': best_iou, }, f'checkpoints/ckpt_epoch{epoch:03d}_iou{best_iou:.4f}.pth')恢复训练时必须按相同顺序加载:
ckpt = torch.load('checkpoints/ckpt_epoch020_iou0.6123.pth') model.load_state_dict(ckpt['model_state_dict']) optimizer.load_state_dict(ckpt['optimizer_state_dict']) scheduler.load_state_dict(ckpt['scheduler_state_dict']) start_epoch = ckpt['epoch'] + 1有两个容易翻车的点。一是模型如果被DataParallel包过,保存的state_dict的键名带module.前缀,恢复训练时要么继续用DataParallel要么手动去掉前缀再加载;二是换不同版本的PyTorch加载旧checkpoint,偶尔会遇到兼容问题,最好在README里记录训练时的PyTorch版本。
提示:写论文时如果改过
num_classes或base_channels,旧权重文件直接加载会报shape不匹配。我曾因为调整类别数忘了重新训练,把旧权重硬加载到一个结构不同的模型上,训练出来的结果乱成一团,浪费了整周时间。结构变了就从头训练,没有捷径。
6. 让baseline变成创新点:模型改进、训练技巧与验证习惯
6.1 给UNet加注意力模块:最小的改动换取可解释的改进点
毕业设计要拿到“对UNet进行改进”的评价,最简单也最能讲清楚的是加注意力机制。注意力UNet在解码器的每次跳跃连接前插入一个Attention Gate,让网络自适应选择编码器特征图里哪些位置的信息更值得传递。对遥感分割来说,这个模块能有效压制背景区域的特征响应,把模型的学习重心放在建筑物、道路这类前景上。
# Attention Gate:压缩两路特征,学习空间权重 class AttentionGate(nn.Module): def __init__(self, in_channels): super().__init__() self.W_g = nn.Conv2d(in_channels, in_channels // 2, kernel_size=1) self.W_x = nn.Conv2d(in_channels, in_channels // 2, kernel_size=1) self.psi = nn.Conv2d(in_channels // 2, 1, kernel_size=1) self.relu = nn.ReLU(inplace=True) self.sigmoid = nn.Sigmoid() def forward(self, g, x): # g是解码器上采样后的特征,x是编码器跳跃连接的原始特征 gate = self.sigmoid(self.psi(self.relu(self.W_g(g) + self.W_x(x)))) return x * gate # 逐位置加权接入位置在解码器拼接之前,把原始编码器特征替换成注意力加权后的特征。这个改造量很小,论文里能作为“对UNet跳跃连接的改进”这一个独立创新点来写。在常用遥感分割数据集上,单独加Attention Gate通常能给mean IoU带来1到3个点的提升,而且不需要调整其他训练参数。
6.2 深监督与多尺度输入:训练稳定性和精度一起解决
深监督的原理是对UNet解码器的每一层都添加一个1x1卷积输出分割结果,每个层单独计算损失,最终loss是各层损失的加权和。这么做能缓解深层网络的梯度衰减问题。遥感图里不同尺度的地物——大块的植被区域、细长的道路——会被解码器不同深度的层侧重捕捉,深监督让它们从正确位置获得梯度信号。
多尺度输入的思路更直接:训练时将256x256的patch以0.5、0.75、1.0三种比例分别输入模型,特征层取加权平均。如果不想做大改动,一个折中方案是在编码器最底部加膨胀卷积金字塔。这个技巧比较适合实验对比,在论文中作为“消融实验”的环节,不改网络结构只看多尺度输入的贡献,对于毕业设计来说扎实而且逻辑自洽。
6.3 在自建数据集上验证模型的完整流程:预测结果反向检查数据问题
跑完训练和推理,别急着把图存下来就算完事。我每次跑完一套遥感分割模型之后强制走一遍验证流程:先看每个类别的IoU,而不是只看mean IoU;然后把预测结果叠加到原始图上逐块目检;最后把误分割最严重的预测图找出来,回到原图上确认是标注问题还是模型局限。
实际操作方式是准备一个可视化对比脚本,将原图、标注、预测结果三张图拼接成一行输出。如果某个类别的误检总是集中出现在某种地形或光照条件下,比如树荫下的道路总是被预测成建筑,这往往说明模型学到了错误的纹理特征,需要对应的增加此类负样本,而不是盲目调参数。
从现在开始,我每次完成UNet类的训练任务,都会强制做一遍“按类目看指标、叠加图目检、回查异常样本”这三个步骤,省下过很多“看似一切正常、实则预测结果没眼见”的时间,希望这些踩坑记录能帮你把毕设推进得更顺利,也能让你少走一些弯路。
如果你已经读到这个位置,说明你已经准备开始实际动手了。压缩包里的create_dataset.ipynb、train.ipynb和predict.ipynb正好对应本文的每一个大步骤。建议按这个顺序走一遍:先跑通默认数据和参数,确认环境无问题,再替换成你自己的遥感数据集。祝顺利。
本文还有配套的精品资源,点击获取