遥感影像分割中类别不均衡？试试用PyTorch实现Focal Loss，我的mIoU提升了5%

遥感影像分割中类别不均衡用PyTorch实现Focal Loss提升5% mIoU的实战指南在遥感影像分割任务中建筑、道路、水体等类别的像素分布往往极不均衡——道路可能只占图像的5%而植被覆盖可能超过40%。这种不平衡会导致模型偏向于学习多数类特征而忽视少数类。传统解决方案如加权交叉熵虽然有效但在极端不平衡场景下仍显不足。本文将分享如何通过Focal Loss解决这一痛点并附上可复现的PyTorch实现代码。1. 为什么传统方法在遥感分割中失效遥感影像的特殊性在于其地物分布具有显著的空间异质性。以某卫星图像数据集为例我们统计了各类别像素占比类别像素占比传统交叉熵准确率加权交叉熵准确率建筑8.2%63%71%道路4.7%55%68%水体12.1%82%85%植被42.3%95%93%加权交叉熵通过反向调整类别权重确实提升了少数类表现但存在两个致命缺陷静态权重无法区分样本难度对简单样本和困难样本施加相同的权重调整超参数敏感权重系数需要反复调整不同数据集泛化性差实际项目中我们发现当道路类别的权重超过3倍时模型开始出现明显的过拟合现象验证集指标不升反降。2. Focal Loss的遥感适配改造Focal Loss的核心创新在于引入动态权重机制其公式为FL(p_t) -α_t * (1 - p_t)^γ * log(p_t)其中关键参数对遥感任务的影响γ (gamma)控制难易样本权重差异γ0 退化为加权交叉熵γ2 时模型对困难样本的关注度是简单样本的25倍α (alpha)类别平衡系数遥感场景建议初始值设为类别占比的倒数归一化通过实验对比不同参数组合在验证集的表现γ值α策略道路mIoU建筑mIoU整体mIoU1.0固定0.2568.272.178.32.0按类别频率调整73.576.881.73.0动态衰减71.275.380.13. PyTorch实现与工程细节针对遥感影像特点我们改进的标准Focal Loss实现包含以下关键点class AdaptiveFocalLoss(nn.Module): def __init__(self, gamma2, alphaadaptive): super().__init__() self.gamma gamma self.alpha_mode alpha def forward(self, inputs, targets): # 计算原始交叉熵 ce_loss F.cross_entropy(inputs, targets, reductionnone) # 动态计算alpha if self.alpha_mode adaptive: class_counts torch.bincount(targets.flatten()) alpha 1.0 / (class_counts 1e-5) alpha alpha / alpha.sum() alpha alpha[targets] else: alpha 0.25 # 计算概率pt pt torch.exp(-ce_loss) fl_loss alpha * (1 - pt)**self.gamma * ce_loss return fl_loss.mean()工程实践中需要注意概率稳定性处理添加torch.exp(-ce_loss)而非直接使用softmax输出避免数值溢出批量统计策略小批量训练时建议使用滑动平均维护类别频率多尺度融合对深层特征施加更强的Focal权重4. 何时不该使用Focal Loss尽管Focal Loss在多数场景表现优异但在以下情况可能适得其反高分辨率影像如0.5m/pixel中建筑物边缘清晰易识别类别间特征差异显著如水体与植被的NDVI值明显可分标注噪声较大时会过度强化错误样本的权重某城市建筑物提取项目中不同方法的对比结果方法训练时间mIoU小目标召回率标准交叉熵2.1h74.2%62%加权交叉熵2.3h76.8%71%Focal Loss (γ2)2.5h81.3%83%Focal Loss CRF3.2h82.7%85%5. 进阶技巧与其他模块的协同优化单纯的损失函数改进可能遇到性能瓶颈我们推荐组合策略样本预处理对少数类进行适度过采样使用Copy-Paste数据增强架构层面# 在UNet跳跃连接处添加类别感知注意力 class ClassAwareAttention(nn.Module): def __init__(self, in_channels): super().__init__() self.query nn.Conv2d(in_channels, in_channels//8, 1) self.key nn.Conv2d(in_channels, in_channels//8, 1) self.value nn.Conv2d(in_channels, in_channels, 1) def forward(self, x, class_mask): # class_mask为低分辨率类别预测 B, C, H, W x.shape q self.query(x).view(B, -1, H*W) k self.key(class_mask).view(B, -1, H*W) v self.value(x).view(B, -1, H*W) attn torch.softmax(q k.transpose(1,2), dim-1) return (attn v).view(B, C, H, W)后处理优化对Focal Loss预测结果进行形态学处理结合DSM高程数据过滤误检在实际部署中发现当配合使用频域增强策略时Focal Loss对细小道路的检测率可再提升2-3个百分点。这种提升在夜间遥感影像中尤为明显因为道路与背景的对比度通常会降低。