PyTorch遥感影像语义分割实战：Unet架构与组合损失函数优化-编程实验室

简介：语义分割是计算机视觉的核心任务之一，旨在为图像中的每个像素分配一个类别标签，实现像素级的场景理解。其原理在于通过编码器-解码器网络结构（如Unet）提取多尺度特征，并融合上下文信息进行精准分类。这项技术的价值在于能够自动化、精细化地解析图像内容，极大地提升了图像分析的效率和精度。在遥感影像分析、自动驾驶、医疗影像诊断等场景中，语义分割技术发挥着关键作用。针对遥感影像地物分割中普遍存在的类别不平衡问题，组合损失函数（如SoftCrossEntropyLoss与DiceLoss）成为有效的优化策略。本文聚焦于利用PyTorch框架，结合Unet架构与组合损失函数，解决高分辨率遥感影像语义分割的实际工程挑战，为相关应用提供实践参考。

1. 项目概述：从像素到地物，遥感影像的智能“翻译官”

拿到这个项目标题，第一反应是“活儿挺全”。一个基于PyTorch的遥感影像语义分割项目，核心是Unet架构，但亮点在于损失函数的组合拳——SoftCrossEntropyLoss、DiceLoss和LovaszLoss。这可不是简单的“Hello World”级图像分割，而是直接瞄准了高分辨率遥感影像这个硬骨头。遥感影像是什么？简单说，就是卫星或飞机从天上拍下来的大地照片。但它的“高分辨率”意味着海量数据（动不动就是成千上万个像素点），以及极其复杂的地物信息（建筑、道路、植被、水体、裸土等交织在一起）。传统的目视解译费时费力，而我们的目标，就是训练一个AI模型，让它像一位经验丰富的解译员，能自动、精准地从影像中“抠”出不同的地物类别，为每一像素打上标签，这就是语义分割。

为什么是Unet？在医学影像分割领域一战成名的Unet，其编码器-解码器结构加跳跃连接的设计，天生就是为了捕获多尺度上下文信息并实现精准的像素级定位，这正好契合了遥感影像中地物目标尺度多变、边界复杂的特点。而标题中提到的损失函数组合，更是直指遥感分割的核心痛点：类别不平衡。一片城区影像里，建筑和道路可能只占20%的像素，而背景（如植被、阴影）占了80%。如果只用普通的交叉熵损失，模型会倾向于把所有像素都预测为背景来获得一个“不错”的整体准确率，但对小目标（如车辆、单棵树）的识别就会一塌糊涂。DiceLoss直接优化分割区域的重叠度，对小目标更敏感；LovaszLoss则是基于子模优化的“软”版IoU损失，在理论上有更好的性质；而SoftCrossEntropyLoss（通常指带标签平滑的交叉熵）则能缓解过拟合，提升模型泛化能力。这个项目，本质上就是在用最前沿的深度学习工具，解决遥感领域最实际、最棘手的问题。

2. 核心需求解析与方案设计思路

2.1 遥感影像智能解译的核心挑战

在动手写代码之前，我们必须先搞清楚我们要解决什么问题。高分辨率遥感影像的语义分割，远非在ImageNet上分类猫狗那么简单，它有几个鲜明的特点：

巨大的图像尺寸：遥感影像单张尺寸通常在512x512到上万像素不等。直接输入网络会爆显存，必须进行裁剪（Crop）或下采样。但下采样会丢失细节，影响小目标识别，因此裁剪成固定大小的小块（如256x256, 512x512）进行训练是主流做法，但这引入了新的问题——上下文信息可能被割裂。
严重的类别不平衡：如前所述，地物分布极不均衡。城市中建筑密集，乡村则农田、林地广布。这要求损失函数必须对少数类别有足够的“注意力”。
复杂的光谱与空间特征：遥感影像通常包含多个波段（如RGB、近红外）。不同地物在不同波段下的反射特性不同，这既是挑战也是机遇。我们的模型需要能有效融合多光谱信息。
模糊的边界与同物异谱/同谱异物：农田和草地的边界可能很模糊；同为“水体”，清澈的湖泊和浑浊的河流在光谱上差异很大；而“裸土”和某些材质的“屋顶”可能光谱相似。这要求模型具备强大的特征学习和上下文理解能力。

2.2 技术方案选型：为什么是Unet+组合损失？

面对这些挑战，我们的技术选型思路如下：

骨架网络（Backbone）：标题指定了Unet。这是一个非常稳妥且强大的选择。其对称的编码器-解码器结构，通过跳跃连接将浅层的高分辨率细节特征与深层的抽象语义特征融合，完美解决了定位与分类的矛盾。编码器部分（下采样）我们通常选用在ImageNet上预训练过的网络，如ResNet、EfficientNet或VGG，以加速收敛并提升特征提取能力。解码器部分（上采样）则通过转置卷积或插值+卷积的方式逐步恢复空间分辨率。
损失函数（Loss Function）：这是本项目的精华所在。单一损失函数很难应对所有情况。
- CrossEntropy Loss (CE)：分类任务的基础，衡量预测概率分布与真实标签分布的差异。但它平等对待所有像素，在类别不平衡时，会被大类别主导。
- SoftCrossEntropyLoss：可以理解为对标准CE的改进，通常通过标签平滑（Label Smoothing）实现。它将硬标签（如[0, 0, 1]）转化为软标签（如[0.1, 0.1, 0.8]），给非目标类别一个很小的概率，从而减轻模型对训练数据的过拟合，增加泛化能力，对噪声标签也有一定的鲁棒性。在遥感数据标注质量参差不齐的情况下，这是一个实用的技巧。
- Dice Loss：源于医学影像分割，直接优化预测区域与真实区域的交集与并集之比（Dice系数）。它对区域面积敏感，能有效缓解类别不平衡问题，因为它的计算是逐类进行的，小类别的预测错误会带来显著的损失值上升。公式为：Dice Loss = 1 - (2*|X∩Y| + ε) / (|X|+|Y| + ε)，其中ε为平滑项防止除零。
- Lovasz Loss：一种基于Lovasz扩展的、直接优化IoU（交并比）的替代损失函数。与Dice Loss类似，它也是为分割任务设计的，且是子模的，具有更好的理论性质。在实际应用中，Lovasz Loss对于边界优化和提升小目标IoU往往有不错的效果，但计算相对复杂一些。

我们的策略是组合使用这些损失函数，例如：Total Loss = α * SoftCE + β * DiceLoss + γ * LovaszLoss。通过调整权重α, β, γ，我们可以让模型同时关注像素级的分类准确性、区域级的重叠度以及边界优化。这是一种经验性的“调参”，但背后有明确的优化目标。

评估指标：不能只看整体准确率（Accuracy）。在类别不平衡下，Accuracy是虚高的。我们必须关注平均交并比（mIoU）——这是语义分割的金标准。此外，F1-Score（尤其是各类别的F1）、平均精度（mPA）也是重要的参考。

3. 环境搭建与数据准备实操

3.1 PyTorch深度学习环境配置要点

工欲善其事，必先利其器。一个稳定、高效的开发环境是项目成功的基石。这里我强烈建议使用Anaconda进行Python环境管理，它能完美解决不同项目间依赖包版本冲突的问题。

创建并激活虚拟环境：

conda create -n rs_seg python=3.8 # 推荐Python 3.8，兼容性好 conda activate rs_seg

安装PyTorch：这是核心步骤，务必与你的CUDA版本匹配。前往 PyTorch官网，利用其提供的安装命令生成器。
- 查看CUDA版本：在命令行输入nvidia-smi，右上角显示的就是你的驱动支持的CUDA最高版本。假设你的是11.7。
- 生成安装命令：在官网选择PyTorch 2.0+（或稳定版本）、你的操作系统、包管理器（Conda或pip）、CUDA 11.7。你会得到类似下面的命令：
```
# 示例，请以官网生成为准 pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu117
```
- 无GPU或Mac用户：选择CPU版本即可。
- 验证安装：在Python中运行import torch; print(torch.__version__); print(torch.cuda.is_available())，应能正确打印版本号并返回True（GPU版）。

安装其他依赖：

pip install opencv-python pillow matplotlib scikit-learn scikit-image tqdm tensorboard # 安装用于计算Lovasz Loss的包（如果不用可跳过） pip install lovasz-losses

注意：网络上很多教程会一股脑地安装torch和torchvision，但版本不匹配是后续各种诡异错误的根源。务必严格按照官网命令安装。如果你的环境之前装过其他版本的PyTorch，最好先conda remove pytorch torchvision清理干净。

3.2 遥感数据集处理与增强策略

数据是模型的“粮食”。遥感领域常用的公开数据集有ISPRS Vaihingen/Potsdam、DeepGlobe Land Cover、LoveDA等。这里以处理一个通用的遥感数据集为例，讲解关键步骤。

数据目录结构：建议采用如下清晰的结构，方便Dataset类读取。

dataset/ ├── images/ # 存放原始影像（.tif, .png, .jpg） │ ├── train/ │ └── val/ ├── masks/ # 存放对应的标签图（单通道，像素值为类别索引） │ ├── train/ │ └── val/ └── classes.csv # （可选）类别名称与索引对应关系

自定义Dataset类：这是PyTorch数据加载的核心。我们需要实现__len__和__getitem__方法。

import os from PIL import Image import torch from torch.utils.data import Dataset import numpy as np class RemoteSensingDataset(Dataset): def __init__(self, image_dir, mask_dir, transform=None): self.image_dir = image_dir self.mask_dir = mask_dir self.transform = transform self.images = sorted(os.listdir(image_dir)) # 确保图像和标签文件名对应 def __len__(self): return len(self.images) def __getitem__(self, idx): img_path = os.path.join(self.image_dir, self.images[idx]) mask_path = os.path.join(self.mask_dir, self.images[idx].replace('.jpg', '_mask.png')) # 根据实际命名调整 image = np.array(Image.open(img_path).convert('RGB')) mask = np.array(Image.open(mask_path).convert('L'), dtype=np.uint8) # 单通道灰度图 # 非常重要：将mask中的像素值转换为类别索引。假设你的mask是RGB伪彩色，需要映射。 # 这里假设mask已经是单通道的类别索引图（0,1,2,...） # 如果mask是RGB，需要写一个颜色到索引的映射函数 color2index(mask) if self.transform: augmented = self.transform(image=image, mask=mask) image = augmented['image'] mask = augmented['mask'] # 将numpy数组转为Tensor image = torch.from_numpy(image).permute(2, 0, 1).float() / 255.0 # [C, H, W], 归一化 mask = torch.from_numpy(mask).long() # [H, W]， 必须是long类型 return image, mask

数据增强（Data Augmentation）：对于数据量有限的遥感任务，增强是提升模型泛化能力的关键。我们使用albumentations库，它支持对图像和掩码进行同步变换。

import albumentations as A from albumentations.pytorch import ToTensorV2 # 训练集增强：强增强 train_transform = A.Compose([ A.RandomResizedCrop(height=256, width=256, scale=(0.5, 1.0)), # 随机裁剪缩放 A.HorizontalFlip(p=0.5), A.VerticalFlip(p=0.5), A.RandomRotate90(p=0.5), A.OneOf([ # 颜色抖动 A.RandomBrightnessContrast(p=1), A.RandomGamma(p=1), A.HueSaturationValue(p=1), ], p=0.5), A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), # ImageNet统计量，用预训练模型时建议 ToTensorV2(), ]) # 验证集增强：仅做归一化和尺寸调整（或中心裁剪） val_transform = A.Compose([ A.Resize(height=256, width=256), A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ToTensorV2(), ])

实操心得：遥感影像增强要符合物理事实。例如，随机旋转是合理的（卫星视角可能变化），但过度的弹性扭曲可能不合适。RandomResizedCrop能模拟不同尺度的地物，非常有效。归一化时使用ImageNet的均值和标准差，是因为我们的Backbone是在ImageNet上预训练的，这有助于稳定训练。

4. 模型构建：Unet架构与损失函数实现

4.1 构建灵活的Unet模型

我们将实现一个支持不同预训练Backbone的Unet。这里以torchvision.models中的resnet34为例。

import torch import torch.nn as nn from torchvision import models import torch.nn.functional as F class UNet(nn.Module): def __init__(self, n_channels=3, n_classes=6, bilinear=True): super(UNet, self).__init__() self.n_channels = n_channels self.n_classes = n_classes self.bilinear = bilinear # 使用预训练的ResNet作为编码器，只取到layer4之前的部分，丢弃全连接层 backbone = models.resnet34(pretrained=True) self.inc = nn.Sequential(backbone.conv1, backbone.bn1, backbone.relu, backbone.maxpool) self.down1 = backbone.layer1 # 输出通道 64 self.down2 = backbone.layer2 # 输出通道 128 self.down3 = backbone.layer3 # 输出通道 256 self.down4 = backbone.layer4 # 输出通道 512 # 解码器部分 factor = 2 if bilinear else 1 self.up1 = Up(512, 256 // factor, bilinear) self.up2 = Up(256, 128 // factor, bilinear) self.up3 = Up(128, 64 // factor, bilinear) self.up4 = Up(64, 64, bilinear) # 这里输出64通道 self.outc = OutConv(64, n_classes) def forward(self, x): x1 = self.inc(x) # [B, 64, H/2, W/2] x2 = self.down1(x1) # [B, 64, H/4, W/4] x3 = self.down2(x2) # [B, 128, H/8, W/8] x4 = self.down3(x3) # [B, 256, H/16, W/16] x5 = self.down4(x4) # [B, 512, H/32, W/32] x = self.up1(x5, x4) x = self.up2(x, x3) x = self.up3(x, x2) x = self.up4(x, x1) logits = self.outc(x) # [B, n_classes, H, W] return logits # 定义上采样模块 class Up(nn.Module): def __init__(self, in_channels, out_channels, bilinear=True): super().__init__() if bilinear: self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True) self.conv = DoubleConv(in_channels, out_channels, in_channels // 2) else: self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2) self.conv = DoubleConv(in_channels, out_channels) def forward(self, x1, x2): # x1: 来自上一层的特征，空间尺寸小，通道多 # x2: 来自编码器对应层的特征，空间尺寸大，通道少 x1 = self.up(x1) # 处理尺寸可能不匹配的问题（由于下采样取整等） 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]) x = torch.cat([x2, x1], dim=1) # 在通道维度拼接 return self.conv(x) # 定义双卷积模块和输出卷积模块（略，为标准实现） class DoubleConv(nn.Module): def __init__(self, in_channels, out_channels, mid_channels=None): super().__init__() # ... 两个(Conv2d -> BN -> ReLU)的组合 def forward(self, x): # ... class OutConv(nn.Module): def __init__(self, in_channels, out_channels): super().__init__() self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1) def forward(self, x): return self.conv(x)

注意事项：使用预训练Backbone时，第一层卷积的输入通道数默认为3（RGB）。如果你的影像是多波段（如4波段RGBN），需要修改第一层卷积的权重。通常的做法是复制RGB通道的权重均值到新增的通道，或随机初始化新增通道的权重。

4.2 组合损失函数的代码实现

现在，我们来实现标题中提到的“SoftCrossEntropyLoss与DiceLoss或LovaszLoss”的组合。

import torch import torch.nn as nn import torch.nn.functional as F import numpy as np class SoftCrossEntropyLoss(nn.Module): """带标签平滑的交叉熵损失""" def __init__(self, smoothing=0.1): super(SoftCrossEntropyLoss, self).__init__() self.smoothing = smoothing def forward(self, pred, target): # pred: [B, C, H, W] 网络输出的logits（未经过softmax） # target: [B, H, W] 类别索引 log_probs = F.log_softmax(pred, dim=1) # 计算log概率 n_classes = pred.size(1) # 将target转换为one-hot，并进行标签平滑 with torch.no_grad(): target_one_hot = torch.zeros_like(pred).scatter_(1, target.unsqueeze(1), 1) target_one_hot = target_one_hot * (1 - self.smoothing) + self.smoothing / n_classes # 计算损失 loss = - (target_one_hot * log_probs).sum(dim=1).mean() return loss class DiceLoss(nn.Module): """Dice Loss，支持多类别，处理类别不平衡""" def __init__(self, weight=None, ignore_index=-100, eps=1e-7): super(DiceLoss, self).__init__() self.weight = weight # 可选的类别权重 self.ignore_index = ignore_index self.eps = eps def forward(self, pred, target): # pred: [B, C, H, W] # target: [B, H, W] if pred.dim() != 4: raise ValueError(f'Expected 4D tensor (B,C,H,W), got {pred.dim()}D') if target.dim() != 3: raise ValueError(f'Expected 3D tensor (B,H,W), got {target.dim()}D') n_classes = pred.size(1) # 将pred通过softmax并忽略ignore_index的标签 probs = F.softmax(pred, dim=1) target_one_hot = F.one_hot(target.clamp(0), n_classes).permute(0, 3, 1, 2).float() # [B, C, H, W] if self.ignore_index is not None and self.ignore_index >= 0: mask = (target != self.ignore_index).unsqueeze(1).expand_as(probs) probs = probs * mask target_one_hot = target_one_hot * mask dims = (0, 2, 3) # 对Batch, Height, Width维度求和，保留类别维度 intersection = torch.sum(probs * target_one_hot, dim=dims) cardinality = torch.sum(probs + target_one_hot, dim=dims) dice_score = (2. * intersection + self.eps) / (cardinality + self.eps) dice_loss = 1. - dice_score if self.weight is not None: dice_loss = dice_loss * self.weight return dice_loss.mean() class CombinedLoss(nn.Module): """组合损失：SoftCE + Dice + Lovasz (可选)""" def __init__(self, ce_weight=1.0, dice_weight=1.0, lovasz_weight=0.0, smooth=0.1, ignore_index=-100): super(CombinedLoss, self).__init__() self.ce_weight = ce_weight self.dice_weight = dice_weight self.lovasz_weight = lovasz_weight self.soft_ce = SoftCrossEntropyLoss(smoothing=smooth) self.dice = DiceLoss(ignore_index=ignore_index) if lovasz_weight > 0: try: from lovasz_losses import lovasz_softmax self.lovasz = lovasz_softmax except ImportError: print("Warning: lovasz-losses not installed. Lovasz loss will be ignored.") self.lovasz_weight = 0.0 self.lovasz = None else: self.lovasz = None def forward(self, pred, target): loss = 0.0 if self.ce_weight > 0: loss += self.ce_weight * self.soft_ce(pred, target) if self.dice_weight > 0: loss += self.dice_weight * self.dice(pred, target) if self.lovasz_weight > 0 and self.lovasz is not None: # lovasz_softmax 期望输入为 [B, C, H, W] 的logits和 [B, H, W] 的标签 lovasz_loss = self.lovasz(F.softmax(pred, dim=1), target, ignore=self.ignore_index) loss += self.lovasz_weight * lovasz_loss return loss

实操心得：损失函数权重的设置是门艺术。通常可以从[1.0, 1.0, 0.5]（CE, Dice, Lovasz）开始尝试。如果数据集类别极度不平衡，可以适当提高DiceLoss的权重（如1.5或2.0）。Lovasz Loss对边界优化效果好，但训练初期可能不稳定，可以设置一个较小的权重（如0.2）或在训练后期加入。务必在验证集上监控各类别的IoU变化来调整权重。

5. 模型训练、验证与调优全流程

5.1 训练循环与验证策略

有了数据、模型和损失函数，我们就可以组装训练流程了。这里的关键是写好训练和验证的循环，并集成TensorBoard进行可视化。

import torch.optim as optim from torch.utils.data import DataLoader from torch.utils.tensorboard import SummaryWriter import time def train_epoch(model, dataloader, criterion, optimizer, device, epoch, writer=None): model.train() running_loss = 0.0 for i, (images, masks) in enumerate(dataloader): images, masks = images.to(device), masks.to(device) optimizer.zero_grad() outputs = model(images) # [B, C, H, W] loss = criterion(outputs, masks) loss.backward() optimizer.step() running_loss += loss.item() if i % 10 == 9: # 每10个batch打印一次 print(f'Epoch [{epoch}], Step [{i+1}/{len(dataloader)}], Loss: {loss.item():.4f}') if writer: writer.add_scalar('Training Loss/step', loss.item(), epoch * len(dataloader) + i) epoch_loss = running_loss / len(dataloader) return epoch_loss def validate(model, dataloader, criterion, device, num_classes): model.eval() val_loss = 0.0 conf_matrix = torch.zeros(num_classes, num_classes) # 用于计算mIoU的混淆矩阵 with torch.no_grad(): for images, masks in dataloader: images, masks = images.to(device), masks.to(device) outputs = model(images) loss = criterion(outputs, masks) val_loss += loss.item() # 计算预测结果，用于评估指标 preds = torch.argmax(outputs, dim=1) # [B, H, W] for t, p in zip(masks.view(-1), preds.view(-1)): conf_matrix[t.long(), p.long()] += 1 val_loss /= len(dataloader) # 计算mIoU intersection = torch.diag(conf_matrix) union = conf_matrix.sum(dim=1) + conf_matrix.sum(dim=0) - intersection iou = intersection / (union + 1e-8) miou = iou.mean().item() return val_loss, miou, iou def main(): device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = UNet(n_channels=3, n_classes=6).to(device) criterion = CombinedLoss(ce_weight=1.0, dice_weight=1.0, lovasz_weight=0.2) optimizer = optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-4) scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='max', factor=0.5, patience=5, verbose=True) train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True, num_workers=4, pin_memory=True) val_loader = DataLoader(val_dataset, batch_size=4, shuffle=False, num_workers=4, pin_memory=True) writer = SummaryWriter('runs/remote_sensing_exp1') best_miou = 0.0 for epoch in range(100): print(f'\nEpoch {epoch+1}/100') train_loss = train_epoch(model, train_loader, criterion, optimizer, device, epoch, writer) val_loss, miou, per_class_iou = validate(model, val_loader, criterion, device, num_classes=6) print(f'Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}, mIoU: {miou:.4f}') writer.add_scalar('Loss/train', train_loss, epoch) writer.add_scalar('Loss/val', val_loss, epoch) writer.add_scalar('Metrics/mIoU', miou, epoch) # 学习率调度 scheduler.step(miou) # 保存最佳模型 if miou > best_miou: best_miou = miou torch.save({ 'epoch': epoch, 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'best_miou': best_miou, }, 'best_model.pth') print(f'Best model saved with mIoU: {best_miou:.4f}') writer.close() if __name__ == '__main__': main()

5.2 超参数调优与训练技巧

训练深度学习模型就像炼丹，火候（超参数）很重要。

学习率（LR）：这是最重要的参数。对于AdamW优化器，1e-4是一个不错的起点。可以使用ReduceLROnPlateau调度器，当验证集指标（如mIoU）不再提升时自动降低学习率。
批量大小（Batch Size）：在显存允许的情况下尽可能大。大的Batch Size能使梯度估计更稳定，但可能降低模型泛化能力。通常从8或16开始尝试。
优化器：AdamW（Adam with decoupled weight decay）是目前图像任务的主流选择，比标准的Adam泛化性能更好。权重衰减（weight_decay）通常设为1e-4。
数据增强强度：增强太弱，模型容易过拟合；增强太强，模型学不到有效特征。需要根据数据集大小调整。小数据集需要更强的增强。
损失函数权重：这是一个需要反复实验的过程。一个实用的策略是：先只用SoftCE训练几个epoch让模型初步收敛，然后加入DiceLoss，观察小类别IoU是否提升。最后在训练中后期加入LovaszLoss来精细化边界。
类别权重：如果某些类别（如“汽车”）的样本特别少，可以在DiceLoss或SoftCE中传入类别权重（weight参数），给少数类别更高的惩罚。权重可以设置为1 / sqrt(class_frequency)或median_frequency / class_frequency。

踩坑记录：我曾在一个项目中，验证集Loss持续下降但mIoU不升反降。排查后发现是数据泄露——验证集的图像和训练集来自同一张大幅影像的不同裁剪块，且地物分布高度相似。这导致模型在验证集上“记住”了特征，而非真正泛化。解决方案是确保训练集和验证集在空间上完全独立（例如，用不同城市的影像，或同一城市但距离很远的区域）。

6. 模型推理、可视化与性能分析

6.1 单张影像推理与后处理

训练好的模型最终要用于预测新的影像。推理过程需要注意尺寸匹配和结果后处理。

def predict_single_image(model, image_path, transform, device, original_size=None): """对单张影像进行预测""" model.eval() # 1. 读取和预处理图像 image = Image.open(image_path).convert('RGB') if original_size is None: original_size = image.size # (W, H) input_image = transform(image=np.array(image))['image'] # 应用验证时的transform input_tensor = input_image.unsqueeze(0).to(device) # [1, C, H, W] # 2. 推理 with torch.no_grad(): output = model(input_tensor) # [1, C, H, W] pred_mask = torch.argmax(output, dim=1).squeeze().cpu().numpy() # [H, W] # 3. 将预测mask缩放到原始影像尺寸 pred_mask_resized = cv2.resize(pred_mask.astype(np.uint8), (original_size[0], original_size[1]), interpolation=cv2.INTER_NEAREST) # 必须用最近邻，保持类别标签 return pred_mask_resized def overlay_visualization(original_image, pred_mask, color_map): """将预测结果以半透明颜色叠加在原图上可视化""" # color_map: 字典，{类别索引: (R, G, B)} h, w = pred_mask.shape color_mask = np.zeros((h, w, 3), dtype=np.uint8) for class_idx, color in color_map.items(): color_mask[pred_mask == class_idx] = color # 将彩色mask叠加到原图 overlay = cv2.addWeighted(original_image, 0.6, color_mask, 0.4, 0) return overlay

6.2 模型性能深度分析与错误排查

模型训练完成后，不能只看一个mIoU数字就了事，必须进行细致的分析。

混淆矩阵分析：计算每个类别的精确率（Precision）、召回率（Recall）和F1-Score。这能告诉你模型在哪些类别上混淆严重。例如，如果“道路”和“人行道”经常分错，说明它们的特征在数据中可能太相似，需要考虑增加更多区分性的训练样本或调整模型特征。
可视化错误案例：在TensorBoard或单独保存一些预测结果最差的样本。仔细观察：
- 边界模糊：是否是Lovasz Loss权重不够？
- 小目标漏检：是否是Dice Loss权重不够或训练数据中小目标样本太少？
- 整块错分：是否是类别间光谱特征过于相似？是否需要引入多时相或高程信息？
推理速度测试：在目标部署硬件（如服务器GPU或边缘设备）上测试模型的FPS（帧每秒）。如果速度不达标，需要考虑模型轻量化，如：
- 将Backbone替换为MobileNetV3、EfficientNet-Lite等轻量网络。
- 使用深度可分离卷积（Depthwise Separable Convolution）改造Unet（即标题热词中的“深度可分离卷积unet”）。
- 进行模型剪枝或量化。

7. 项目总结与进阶方向

走完整个流程，你会发现，基于Unet的遥感影像分割是一个经典但充满细节的工程。它不仅仅是把模型跑通，更涉及到数据工程、损失函数设计、训练策略和错误分析等一系列环环相扣的步骤。

我个人在多个类似项目中的体会是，数据质量决定上限，模型和损失函数决定逼近上限的速度。花在数据清洗、标注核对和数据增强策略上的时间，往往比调参带来的收益更大。对于损失函数，没有银弹，SoftCE+Dice的组合在大多数情况下已经能提供一个很强的基线，Lovasz可以作为锦上添花的优化手段。

这个项目还可以向多个方向扩展：