高分遥感图像语义分割实战：从PyTorch环境搭建到UNet模型训练与mIoU优化

简介这份资源面向遥感图像处理方向的研究者、工程师及具备一定深度学习基础的学习者提供基于Pytorch实现高分辨率遥感图像语义分割的完整教程与配套数据集帮助解决地物信息提取中端到端训练与评估的实操问题。压缩包共1029个文件约577.48MB以819个png图像样本、35个py脚本、171个zbak备份文件为主另含csv标注索引、jpg效果图与md说明文档覆盖数据预处理、模型训练到结果评估的完整链路。教程从遥感图像基本概念讲起逐步深入到网络结构选择、优化策略及标注工具制作像素级标签的方法并附数据集下载链接方便学习者亲身体验全流程。目前已有94人学习下载适合希望将语义分割技术落地到环境监测、资源勘探等场景的读者参考。1. 高分遥感图像语义分割从 0.1 米影像里抠出每一类地物手里有一景 0.1 米分辨率的航拍正射影像几万乘几万像素要分出建筑、道路、水体、植被、裸地五类。用传统阈值法调 HSV 阈值调到怀疑人生阴影和深色沥青永远打架。这就是高分遥感语义分割要解决的问题给每个像素打一个类别标签输出一张和原图等大的掩膜。它和普通自然图像分割最大的差别在于——目标尺度极端不均衡一栋楼可能占几万像素一辆车只有几十像素同时地物边界模糊农田和草地肉眼都难分。适合谁做做国土调查、耕地识别、城市扩张监测、灾害评估的工程师以及想用 PyTorch 把语义分割真正跑通一遍的算法同学。这篇不聊虚的从环境、数据、模型到训练排错按能复现的路径走一遍。2. 环境与数据PyTorch 装完只是开始数据集才是分水岭2.1 PyTorch 环境搭建CUDA 版本对不上后面全是玄学遥感分割模型动辄几千万参数输入又是 512×512 起步CPU 训练基本等于放弃。所以第一步是把 GPU 版 PyTorch 装对。我一般用 conda 建独立环境避免和系统里的包打架。先确认显卡驱动支持的 CUDA 上限再去 PyTorch 官网拿对应命令不要凭记忆抄。# 创建独立环境python 版本选 3.10兼容性最好 conda create -n rs_seg python3.10 -y conda activate rs_seg # 安装 GPU 版 pytorchcu121 表示 CUDA 12.1按自己驱动改 # 这条命令来自 pytorch 官方索引不要用 pip install torch 裸装 pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121 # 验证 GPU 是否真的可用这一步必须做 python -c import torch; print(torch.__version__, torch.cuda.is_available(), torch.cuda.get_device_name(0))逻辑说明--index-url指向 PyTorch 官方 wheel 源保证拿到的是编译好 CUDA 的版本而不是 CPU 版。参数说明cu121要和你nvidia-smi右上角显示的 CUDA Version 匹配或更低torch.cuda.is_available()返回 False 就说明装成了 CPU 版或者驱动太旧。常见翻车是 pip 默认源里也有个 torch版本号看着对实际是 CPU 轮子训练时才发现devicecuda报错。装完再补几个遥感分割常用库pip install opencv-python rasterio albumentations tqdm tensorboard。rasterio 用来读带地理坐标的 GeoTIFFalbumentations 做数据增强比 torchvision 的 transforms 更灵活支持同步变换图像和掩膜。2.2 高分遥感数据集怎么选、怎么下标题里带「数据集下载」这是很多人卡住的地方。遥感分割数据集和自然图像不一样标注成本极高公开的没几个。我按用途分三类说。第一类是通用地物分割最常用的是 LoveDA 和 Potsdam。LoveDA 覆盖城市和农村7 类地物0.3 米分辨率适合练手和跑 baseline。Potsdam 是 ISPRS 提供的城市正射影像6 类0.05 米分辨率标注质量高做建筑和道路分割的首选。第二类是专项任务比如耕地识别可以用 CLCD 数据集它本身是变化检测数据但单时相可以拿来当耕地/非耕地二分类。第三类是自己标注用 labelme 或 ArcGIS 画多边形导出成掩膜。下载渠道上LoveDA 和 Potsdam 在各自官方页面提供需要填邮箱申请别去第三方网盘找版本经常被改过。下载后目录结构一般是Train/Urban/images_png和对应masks_png先确认图像和掩膜的像素值范围掩膜是 0/1/2… 的类别索引不是 0/255 的二值图这个搞错训练直接崩。import numpy as np from PIL import Image # 检查一张掩膜的类别分布确认标注格式 mask np.array(Image.open(masks_png/000001.png)) print(shape:, mask.shape, dtype:, mask.dtype) print(unique values:, np.unique(mask)) # 正常应输出 [0 1 2 3 4 5 6]如果只有 [0 255] 说明是二值掩膜需要重映射逻辑说明遥感掩膜常见两种坑一是类别索引从 1 开始0 是背景二是用 255 表示忽略区域。训练前必须统一成 0 到 N-1 的连续整数否则 CrossEntropyLoss 的ignore_index和类别数对不上。参数说明np.unique看实际出现的值如果最大值远大于类别数就要写个映射表转换。2.3 数据增强遥感图像不能随便翻转自然图像增强里随机旋转、翻转随便用遥感图像有个约束——地理方向有意义。垂直翻转会让南北颠倒如果任务和太阳光照方向、地形坡向相关就会引入错误先验。我一般只用水平翻转、90 度整数倍旋转以及颜色抖动。裁剪用随机 crop 到 512×512因为高分影像太大整图塞不进显存。import albumentations as A train_tf A.Compose([ A.RandomCrop(512, 512), # 裁到固定尺寸适配显存 A.HorizontalFlip(p0.5), # 水平翻转安全 A.RandomRotate90(p0.5), # 90 度整数倍旋转不改变地理朝向语义 A.ColorJitter(0.2, 0.2, 0.2, 0.05, p0.3), # 亮度对比度饱和度抖动 ]) # 注意图像和掩膜必须同步变换albumentations 的 Compose 自动处理 aug train_tf(imageimg, maskmask)逻辑说明A.Compose保证 image 和 mask 用同一组随机参数避免图像转了掩膜没转。参数说明RandomCrop的 512 是显存和感受野的折中8G 显存跑 UNet 建议 51212G 以上可以上 768。ColorJitter的四个参数分别是亮度、对比度、饱和度、色调的扰动幅度遥感影像色调变化大色调扰动给 0.05 就够给大了会把水体颜色改得不像水。3. 模型选型UNet、DeepLabV3 还是 SegFormer别盲目追新3.1 三种主流语义分割模型的适用边界遥感分割的模型选型核心看三件事目标尺度跨度、边界精度要求、算力预算。我把常用的三个摆一起对比。模型骨干优势短板适用场景UNet自编码结构简单、小数据也能收敛、边界细感受野有限、大目标易碎建筑提取、样本少DeepLabV3ResNet/Xception空洞卷积多尺度、大目标稳参数多、边界略糊多类地物、大场景SegFormerMiT Transformer全局建模强、多尺度融合好吃数据、显存高大数据集、精度优先我的经验是样本少于 5000 张直接上 UNet别折腾 Transformer数据不够它学不出全局关系反而过拟合。样本上万、类别多、追求 mIoU再考虑 SegFormer。DeepLabV3 是稳妥的中间选择工程上部署成熟。3.2 用 UNet 跑通第一个 baseline先别急着上复杂模型用 UNet 把整条链路跑通确认数据、损失、评估都对再换模型。下面是一个精简的 UNet 实现重点是双卷积块和下采样上采样的拼接。import torch import torch.nn as nn class DoubleConv(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.net 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.net(x) class UNet(nn.Module): def __init__(self, in_ch3, num_classes7): super().__init__() # 编码器通道翻倍尺寸减半 self.d1 DoubleConv(in_ch, 64) self.d2 DoubleConv(64, 128) self.d3 DoubleConv(128, 256) self.d4 DoubleConv(256, 512) self.pool nn.MaxPool2d(2) # 瓶颈层 self.bottleneck DoubleConv(512, 1024) # 解码器上采样后与编码器特征拼接 self.up4 nn.ConvTranspose2d(1024, 512, 2, stride2) self.c4 DoubleConv(1024, 512) self.up3 nn.ConvTranspose2d(512, 256, 2, stride2) self.c3 DoubleConv(512, 256) self.up2 nn.ConvTranspose2d(256, 128, 2, stride2) self.c2 DoubleConv(256, 128) self.up1 nn.ConvTranspose2d(128, 64, 2, stride2) self.c1 DoubleConv(128, 64) self.out nn.Conv2d(64, num_classes, 1) def forward(self, x): d1 self.d1(x) d2 self.d2(self.pool(d1)) d3 self.d3(self.pool(d2)) d4 self.d4(self.pool(d3)) b self.bottleneck(self.pool(d4)) # 拼接时通道数翻倍所以 DoubleConv 输入是 1024 u4 self.c4(torch.cat([self.up4(b), d4], dim1)) u3 self.c3(torch.cat([self.up3(u4), d3], dim1)) u2 self.c2(torch.cat([self.up2(u3), d2], dim1)) u1 self.c1(torch.cat([self.up1(u2), d1], dim1)) return self.out(u1)逻辑说明编码器每层池化降采样扩大感受野解码器用转置卷积上采样再和同层编码器特征torch.cat拼接把丢失的空间细节找回来。参数说明num_classes要和你掩膜的类别数严格一致LoveDA 是 7Potsdam 是 6二分类就是 2。in_ch默认 3如果输入加了 NDVI 等额外波段改成对应通道数。拼接后通道翻倍这点最容易写错c4的输入是 1024 不是 512。3.3 损失函数交叉熵不够要加 Dice遥感分割里类别极不均衡道路、车辆这种小目标像素占比可能不到 1%纯交叉熵会让模型倾向于全预测成背景mIoU 看着还行但小目标全丢。常见做法是交叉熵加 Dice 损失Dice 直接优化预测和真值的重叠度对小目标更敏感。import torch.nn.functional as F def dice_loss(pred, target, num_classes, eps1e-6): # pred: [B, C, H, W] logits, target: [B, H, W] 类别索引 pred F.softmax(pred, dim1) target_onehot F.one_hot(target, num_classes).permute(0, 3, 1, 2).float() # 逐类别算 Dice再平均 inter (pred * target_onehot).sum(dim(0, 2, 3)) union pred.sum(dim(0, 2, 3)) target_onehot.sum(dim(0, 2, 3)) dice (2 * inter eps) / (union eps) return 1 - dice.mean() def total_loss(pred, target, num_classes): ce F.cross_entropy(pred, target, ignore_index255) # 255 为忽略区域 dc dice_loss(pred, target, num_classes) return ce dc # 权重可调一般 1:1 起步逻辑说明one_hot把类别索引转成 one-hot才能和 softmax 后的概率逐像素相乘。参数说明ignore_index255让标注里表示「不确定」的像素不参与损失遥感标注边界常有这种区域。Dice 的eps防止分母为零。两者权重我一般从 1:1 开始如果小目标还是丢把 Dice 权重提到 2。4. 训练与评估mIoU 涨不上去先查这四处4.1 训练循环与学习率策略训练循环本身不复杂关键是学习率调度和验证时机。遥感分割我一般用 AdamW初始学习率 1e-4配合余弦退火。每个 epoch 后在验证集上算 mIoU保存最好的权重。import torch from torch.optim import AdamW from torch.optim.lr_scheduler import CosineAnnealingLR from tqdm import tqdm device torch.device(cuda if torch.cuda.is_available() else cpu) model UNet(in_ch3, num_classes7).to(device) opt AdamW(model.parameters(), lr1e-4, weight_decay1e-4) sched CosineAnnealingLR(opt, T_max50) # 50 个 epoch 周期 for epoch in range(50): model.train() for img, mask in tqdm(train_loader): img, mask img.to(device), mask.to(device) opt.zero_grad() pred model(img) loss total_loss(pred, mask, num_classes7) loss.backward() opt.step() sched.step() # 验证阶段算 mIoU保存最优 miou evaluate(model, val_loader, device, num_classes7) print(fepoch {epoch} mIoU {miou:.4f})逻辑说明zero_grad清梯度必须在 backward 前sched.step()放在 epoch 末尾。参数说明weight_decay1e-4是 AdamW 的常规值遥感数据量不大时能抑制过拟合。T_max设成总 epoch 数余弦退火会从 1e-4 平滑降到接近 0。如果 loss 前几个 epoch 就震荡把学习率降到 5e-5。4.2 mIoU 的计算与混淆矩阵mIoU 是语义分割的标准指标但很多人算错。正确做法是累加整个验证集的混淆矩阵再逐类算 IoU 取平均而不是每个 batch 算完再平均后者会被小 batch 的类别分布带偏。import numpy as np def evaluate(model, loader, device, num_classes): model.eval() conf np.zeros((num_classes, num_classes), dtypenp.int64) with torch.no_grad(): for img, mask in loader: img img.to(device) pred model(img).argmax(dim1).cpu().numpy() mask mask.numpy() # 过滤忽略区域 valid mask ! 255 # 用 bincount 累加混淆矩阵 idx num_classes * mask[valid].astype(int) pred[valid] conf np.bincount(idx, minlengthnum_classes**2).reshape(num_classes, num_classes) # 逐类 IoU inter np.diag(conf) union conf.sum(1) conf.sum(0) - inter iou inter / np.maximum(union, 1) return np.nanmean(iou)逻辑说明混淆矩阵的行是真值、列是预测bincount把每个像素的「真值×类别数预测」编码成一个整数计数比双重循环快得多。参数说明np.maximum(union, 1)防止某类在验证集里没出现导致除零。nanmean忽略掉没出现的类别。如果某类 IoU 一直是 0说明该类样本太少或标注有问题回去查数据。4.3 推理与拼接大图不能整张塞进去高分影像动辄上万像素推理时必须滑窗切块再拼回整图。切块要有重叠否则拼接处会有明显接缝。重叠区域取概率最大或直接覆盖我一般用重叠 1/4 窗口拼接时后写的覆盖先写的。def sliding_inference(model, image, window512, stride384, num_classes7): # image: [C, H, W] tensor model.eval() C, H, W image.shape prob torch.zeros((num_classes, H, W), deviceimage.device) count torch.zeros((H, W), deviceimage.device) with torch.no_grad(): for y in range(0, H, stride): for x in range(0, W, stride): y2, x2 min(y window, H), min(x window, W) y1, x1 max(y2 - window, 0), max(x2 - window, 0) patch image[:, y1:y2, x1:x2].unsqueeze(0) out torch.softmax(model(patch), dim1).squeeze(0) prob[:, y1:y2, x1:x2] out count[y1:y2, x1:x2] 1 return (prob / count.unsqueeze(0)).argmax(dim0)逻辑说明stride小于window就产生重叠重叠区概率累加后除以计数相当于加权平均接缝更平滑。参数说明window512和训练尺寸一致stride384即重叠 128 像素。如果显存够window 可以放大到 768但要注意和训练尺寸差异太大会掉点。5. 避坑与排查这五个坑我全踩过5.1 掩膜类别索引不连续损失直接 NaN现象训练第一个 batch loss 就是 nan或者报Target out of bounds。原因掩膜里类别是 0、1、2、4、5跳过了 3而模型输出 5 类CrossEntropyLoss 遇到索引 4 越界。解决训练前统计所有掩膜的 unique 值写映射表重映射成连续整数映射表存下来推理时反查。5.2 图像和掩膜尺寸对不上报维度错误现象torch.cat或损失计算时报尺寸不匹配。原因数据增强时图像和掩膜用了不同的随机参数或者原图有黑边、掩膜没有。解决统一用 albumentations 的 Compose 同步变换读数据时断言img.shape[:2] mask.shape[:2]不等就跳过或裁剪。5.3 显存溢出batch size 只能设 1现象CUDA out of memory。原因输入尺寸太大或者 UNet 通道数翻倍后中间特征图占满显存。解决先降输入到 256×256 跑通再逐步加用混合精度torch.cuda.amp能省近一半显存实在不够就梯度累积小 batch 多次 backward 再 step。5.4 mIoU 虚高背景类占了大头现象mIoU 0.85 看着不错但小目标全没分出来。原因背景像素占比 90% 以上模型全预测背景也能拿高分。解决看逐类 IoU别只看平均损失里加 Dice 或 Focal评估时单独看道路、车辆这些稀有类的 IoU。5.5 验证集和训练集同源指标虚高现象验证 mIoU 很高换一景新影像就崩。原因训练验证按 patch 随机划分同一景影像的相邻 patch 同时进了训练和验证空间泄漏。解决按整景影像划分训练用 A 景、验证用 B 景才能反映真实泛化能力。6. 进阶技巧用预训练骨干和测试时增强再榨几个点baseline 跑通后想再提点最划算的两招是换预训练骨干和测试时增强。UNet 的编码器可以换成在 ImageNet 上预训练过的 ResNet 或 EfficientNet遥感数据量小的时候预训练权重能带来明显提升。做法是把 UNet 的编码器替换成torchvision的 resnet34取前几层输出作为多尺度特征。import torchvision class ResNetUNet(nn.Module): def __init__(self, num_classes7): super().__init__() backbone torchvision.models.resnet34(weightsIMAGENET1K_V1) # 取 resnet 的四个阶段输出通道分别是 64/64/128/256/512 self.stem nn.Sequential(backbone.conv1, backbone.bn1, backbone.relu) # /2 self.layer1 backbone.layer1 # /4, 64 self.layer2 backbone.layer2 # /8, 128 self.layer3 backbone.layer3 # /16, 256 self.layer4 backbone.layer4 # /32, 512 # 解码器略结构同 UNet输入通道按上面改 self.out nn.Conv2d(64, num_classes, 1) def forward(self, x): x self.stem(x) c1 self.layer1(x) c2 self.layer2(c1) c3 self.layer3(c2) c4 self.layer4(c3) # 后续上采样拼接注意各层通道数 return self.out(c1) # 示意实际要接完整解码器逻辑说明weightsIMAGENET1K_V1加载 ImageNet 预训练权重遥感图像和自然图像分布有差异但底层边缘、纹理特征通用微调后收敛更快。参数说明resnet34 各阶段输出通道是 64、128、256、512解码器拼接时输入通道要对应改。如果显存紧张用 resnet18精度略降但快很多。测试时增强TTA是另一招推理时对图像做水平翻转、多尺度缩放分别预测后把概率平均。代价是推理时间翻几倍但 mIoU 通常能涨 1 到 2 个点对精度敏感的项目值得做。我一般做水平翻转加 0.75、1.0、1.25 三个尺度概率平均后取 argmax。最后说个习惯每次改完数据或模型先在小样本上过一遍确认 loss 能降、掩膜能出图再上全量训练。我早期图省事直接全量跑结果训练两小时才发现掩膜读错了血泪教训。遥感分割没有银弹数据质量决定上限模型只是逼近上限的工具。把数据清洗和标注一致性做扎实比换十个模型都管用。希望帮到你。本文还有配套的精品资源点击获取