AD-HRNet：融合注意力机制与膨胀卷积的遥感语义分割

简介AD-HRNet是一套面向遥感图像语义分割任务的深度学习源码实现以HRNet为骨干网络融合注意力机制与膨胀卷积适配高分辨率遥感影像中地物边界精细分割与多尺度目标识别场景适合从事遥感图像处理、计算机视觉研究的学生及算法工程师阅读与二次开发。压缩包共101个文件大小仅146KB主要包含Python模型与训练脚本、C/CUDA扩展源码、PyTorch工程配置、XML配置及Markdown说明等其中py文件定义网络结构与训练流程cpp/cu/h文件用于编译InPlace ABN等自定义算子其余文档可辅助环境配置与原理理解。目前已有216人学习下载。通过该源码包读者可获得完整的HRNet语义分割PyTorch工程骨架包括模型定义、加载预训练权重、训练与验证流程、InPlace ABN加速模块等便于在此基础上开展针对性改进实验或将注意力与膨胀卷积思想迁移至其他分割任务。1. AD-HRNet 这个名字背后遥感语义分割为什么要同时引入注意力机制和膨胀卷积遥感影像里一栋建筑可能占几十像素一条小路可能只有两三个像素宽而背后的背景往往是成片农田或水域。直接用 HRNet 这种并行高分辨率网络做语义分割细节保留得不错可在高层语义上会感觉“见得不够远”。膨胀卷积能在不做池化的情况下把有效感受野放大注意力机制则把通道和空间上的有用信息显式挑出来。AD-HRNet 可以看作把这两件事同时接进 HRNet 骨架的一种做法典型应用包括地物分类、建筑物提取、水体分割和变化检测。如果你手头有一份“AD-HRNet 用于遥感图像语义分割的结合注意力机制和膨胀卷积的 HRNet 源码.zip”这篇就是从源码思路讲到落地的完整过程。2. 拆开 AD-HRNet 源码注意力模块和膨胀卷积分别嵌在 HRNet 的哪些位置要读懂一份后缀是 .zip 的源码工程首先得知道网络长什么样。AD-HRNet 不是把注意力机制和膨胀卷积随意堆在 HRNet 外面而是有明确的位置和参数语义。这一章先讲清 HRNet 的结构瓶颈再把两处关键改动的实现顺序拆开。2.1 HRNet 不用编码器-解码器结构靠并行分支保细节HRNet 与 ResNet 加 ASPP 这类编码器不同它不先下采样再上采样而是从 stem 之后始终保持一个高分辨率分支同时逐步添加低分辨率分支并在每个 stage 反复做多分辨率融合。这样做的直接收益是分割结果的空间细节不会被上采样过程抹掉道路、河流、田埂这些细长地物在遥感影像里尤其依赖这种结构。但 HRNet 也有一个实际瓶颈高分辨率分支的卷积感受野有限。stage 越浅每个点能看到的区域越小到了 stage4如果只用 3x3 普通卷积有效感受野仍不足以覆盖大块地物的上下文。遥感图像的语义信息强依赖上下文例如涵洞旁边的深色区域到底是水体还是阴影往往要看更大范围。只靠加深网络会增加参数量和显存所以在 HRNet 基础上补一个显式的上下文扩宽模块是常见做法。AD-HRNet 的思路就是用膨胀卷积扩宽感受野用注意力机制重新校准特征。2.2 注意力机制在源码里的典型实现通道-空间协同注意力CBAM在 AD-HRNet 一类工程里最常见的是把 CBAM 插在 HRNet 的每个 stage 输出之后或仅插在最终分割头之前。CBAM 包含通道注意力和空间注意力两个子模块先算 channel weight再算 spatial weight输出就是反射这个权重的特征图。下面这段 PyTorch 实现可以直接对照源码看。import torch import torch.nn as nn class ChannelAttention(nn.Module): def __init__(self, in_planes, ratio16): super().__init__() self.avg_pool nn.AdaptiveAvgPool2d(1) self.max_pool nn.AdaptiveMaxPool2d(1) self.mlp nn.Sequential( nn.Conv2d(in_planes, in_planes // ratio, 1, biasFalse), nn.ReLU(inplaceTrue), nn.Conv2d(in_planes // ratio, in_planes, 1, biasFalse), ) self.sigmoid nn.Sigmoid() def forward(self, x): avg_out self.mlp(self.avg_pool(x)) max_out self.mlp(self.max_pool(x)) return self.sigmoid(avg_out max_out) class SpatialAttention(nn.Module): def __init__(self, kernel_size7): super().__init__() self.conv nn.Conv2d(2, 1, kernel_size, paddingkernel_size // 2, biasFalse) self.sigmoid nn.Sigmoid() def forward(self, x): avg_out torch.mean(x, dim1, keepdimTrue) max_out, _ torch.max(x, dim1, keepdimTrue) out torch.cat([avg_out, max_out], dim1) return self.sigmoid(self.conv(out))这段代码里通道注意力同时用了平均池化和最大池化原因在于平均池化能看到全局响应最大池化能抓住纹理、边缘这类判别性特征。空间注意力则是在通道维度上做压缩再用 7x7 卷积生成空间权重让网络学会“关注哪些像素位置”。在 AD-HRNet 中通道数通常是 32/64/128/256 这样递进ratio 取 16 时不会给模型增加太多参数。2.3 膨胀卷积在源码里的形态rate、padding 与感受野计算膨胀卷积的原理是在卷积核的元素之间插入空洞rate1 时就是普通卷积rate2 时相当于 3x3 覆盖 5x5 范围。PyTorch 里直接通过 dilation 参数控制# 保持输出分辨率不变的 3x3 膨胀卷积 self.dilated_conv nn.Conv2d( in_channels256, out_channels256, kernel_size3, stride1, dilation4, padding4, biasFalse )这里的 padding 必须等于 dilation才能保证输出尺寸与输入一致。有效感受野的计算公式是[ RF_{new} RF_{old} (k-1) \times (dilation - 1) ]其中 k 是卷积核尺寸。rate2 时3x3 卷积的有效核宽是 5rate4 时是 9。在遥感图像里以 512x512 输入为例stage4 的特征大概是 1/32 分辨率感受野每增加 1对应原图就增加 32 像素所以 dilation 取 4 或 6 就能覆盖大目标。需要注意的是如果把多层 dilation 相同的膨胀卷积直接堆叠会出现 gridding 现象也就是卷积核采样点过于稀疏中间区域漏采。常见做法是让 dilation 沿层递增比如 2、4、8或者改用混合膨胀卷积HDC。2.4 源码里的整体数据流从 stem 到分割头怎么看我一般拿到这种源码包不会先逐行读 loss 和 dataset而是先画一条数据流输入张量经过 stem、stage1 到 stage4、注意力模块和膨胀卷积后进入分割头。下面是一个简化到可以对照阅读的 Python 伪代码x self.stem(input) # 1/2 或 1/4 分辨率 x1 self.stage1(x) # 高分辨率分支 x2 self.stage2(x1) # 新增 1/4 分辨率分支 x3 self.stage3(x2) # 新增 1/8 分辨率分支 x4 self.stage4(x3) # 新增 1/16 分辨率分支 # 常见做法把四个分辨率分支上采样到同一尺寸后 concat fused self.fuse(x1, x2, x3, x4) # 注意力模块加在这里或加到每个 stage 内部 fused self.cbam(fused) # 膨胀卷积组在这里扩大感受野 out self.dilated_context(fused) # 最后 1x1 卷积输出每个像素的类别 logits logits self.head(out)这类源码通常会在配置文件中写明哪些层被注意力模块包裹哪些层用了 dilation。你可以在 train.py 里设置--cfg指定 yaml 或 py 配置然后在 backbone 定义处搜索dilation和CBAM两个关键字就能快速定位改动点。如果源码里加了 TensorBoard 或 wandb hook也可以直接看特征图对比注意力加与不加的差异。模块常见插入位置推荐参数主要作用通道注意力stage 输出后ratio16重标定通道空间注意力通道注意力之后kernel_size7聚焦像素位置膨胀卷积stage4 后或 context headdilation4/8放大感受野HRNet stem输入最前stride2初步下采样3. 把 AD-HRNet 源码跑通环境配置、数据准备与训练命令源码能不能跑起来一半看环境一半看数据。遥感图像语义分割的工程和 ImageNet 分类不太一样输入尺寸、类别数、标签格式都直接影响模型配置。这一章按最小可运行路径来搭。3.1 解压源码并创建独立环境拿到这份 AD-HRNet 源码包后先解压到工作目录再用 conda 建独立环境。不要直接在 base 环境里装依赖版本冲突在分割项目里很容易发生。下面是我常用的操作顺序mkdir -p ~/work/adhrnet cd ~/work/adhrnet unzip AD-HRNet*.zip conda create -n adhrnet python3.9 -y conda activate adhrnet pip install torch2.1.0 torchvision0.16.0 --index-url https://download.pytorch.org/whl/cu118 pip install -r requirements.txtrequirements 里一般会有 mmcv、mmsegmentation、timm、einops。如果源码是基于 mmsegmentation 的建议直接用 python 安装匹配本机 CUDA 的 mmcv 预编译包避免源码编译踩坑。torch 版本和 CUDA 版本要匹配装完先运行python -c import torch;print(torch.cuda.is_available())确认 GPU 可用。3.2 准备遥感语义分割数据以 LoveDA 为例遥感语义分割数据和自然图像不太一样标签通常是一张单通道 PNG像素值直接用类别 ID不是 RGB 伪彩色。以 LoveDA 为例数据组织如下loveda ├── train │ ├── img_dir │ │ ├── 000001.png │ │ └── ... │ └── ann_dir │ ├── 000001.png │ └── ... └── val ├── img_dir └── ann_dir在 mmsegmentation 的配置里data_root 指向 loveda 目录img_dir 和 ann_dir 分别对应影像和标签。LoveDA 共有 7 个类别背景不算在内时 num_classes 填 7但如果你的源码包自带头文件把背景也算进去就会变成 8训练前先确认一下。3.3 修改配置文件里的关键参数AD-HRNet 的配置通常是一个 .py 文件里面分为 model、data、optimizer 三段。下面是一个最小训练配置的参数表参数建议值说明num_classes7按数据集实际类别数改in_channels3RGB 输入为 3多光谱按实际波段改img_size512x512遥感影像常用 512 或 1024batch_size8batch8、512x512 输入约占用 11GB 显存optimizer.lr1e-3AdamW 时常用 1e-3 到 1e-4weight_decay0.01结构风险防止过拟合total_epochs80遥感分割通常比 ImageNet 少eval_interval8000 iter验证一轮打点checkpoint_interval4000 iter保存权重的频率如果显存不足优先把 batch_size 降到 4然后同步把学习率按比例调低到 5e-4千万不要只调 batch_size 不调学习率否则初始 loss 会变得很奇怪。3.4 启动训练和观察 loss配置改好后用下面的命令启动训练python tools/train.py configs/adhrnet/adhrnet_loveda.py --work-dir work_dirs/adhrnet_loveda日志会同时输出到终端和 work_dirs/adhrnet_loveda 下的日志文件。训练过程中我一般会开第二个终端跑tail -f看指标tail -f work_dirs/adhrnet_loveda/20250101_102030.log如果 log 里 loss_sem 在 20 个 epoch 内降到 0.1 以下而 val mIoU 还在 30% 左右说明模型在训练集上记住影像但没有学到可泛化的空间结构优先检查数据增强里的翻转、多尺度裁剪是否打开。如果 loss 整体下降但 val mIoU 在某一步突然掉 10 个点很可能是验证集里出现了训练期间没见过的影像覆盖类型比如从城乡结合部换到纯农田区域。4. AD-HRNet 的调参与避坑膨胀卷积 rate、注意力权重和损失函数训练能跑通只是第一步真正花时间的是把 mIoU 从 60% 提到 70% 这个过程。这一章讲膨胀卷积 rate 怎么配、注意力模块放哪、损失函数怎么组合以及最常见的报错怎么排。4.1 膨胀卷积的 rate 怎么选感受野与 gridding遥感图像里地物尺度跨度大膨胀卷积的 rate 不能拍脑袋。我给一个常见用法在 HRNet stage4 输出的 1/16 分辨率特征上设置 3 个并行的膨胀卷积分支dilation 分别为 2、4、8然后 concat类似 ASPP 但去掉了全局池化。这样每个分支的输出分辨率不变却覆盖了不同范围道路、屋顶这类小物体靠 rate2 分支大型水体、农田靠 rate8 分支。self.context_branches nn.ModuleList([ nn.Conv2d(in_ch, out_ch, 3, paddingrate, dilationrate) for rate in [2, 4, 8] ])要注意dilation 数值越大卷积核覆盖的区域越稀疏边界处的响应越不稳定。如果训练数据和推理数据的分辨率不一致比如训练是 512推理是 2048直接套用 rate8 可能在原图上产生周期性的伪影。此时可以把 rate 按分辨率比例缩放或者干脆用全局上下文池化替代过大的 dilation。4.2 注意力机制插在哪CBAM、CA 和 SE 的差异在 AD-HRNet 这类工程里有人用 SE 通道注意力有人用 CBAM 通道-空间协同注意力也有人用 CA 注意力机制。SE 只做通道重标定计算量最小但对位置信息完全不敏感。CBAM 多了一个空间注意力分支能帮助网络记住“哪个像素重要”。CA 则把空间坐标编码进注意力权重对遥感里有大量长条状结构的情况效果更明显。# coordinate attention 的核心思路 # 将特征沿 H 和 W 方向分别池化再合并 pool_h torch.mean(x, dim3, keepdimTrue) # B,C,H,1 pool_w torch.mean(x, dim2, keepdimTrue) # B,C,1,W # 通过 1x1 卷积、非线性激活、再拆分后作为权重我自己的经验是如果分割目标主要是规则地块和建筑物CBAM 足够如果影像里有大量道路、河流这类方向性明显的目标CA 会比 CBAM 带来 1 到 2 个点的 mIoU 提升。但 CA 的额外开销也要计入推理时间。你可以在第 2 章代码里的 CBAM 位置把两个子模块注释掉只保留空间或通道部分分别做一次短训练对比。4.3 遥感长尾下的损失函数组合遥感语义分割最常见的坑是类别不平衡比如“道路”像素占比不到 1%却贡献了交叉熵损失的大部分。单用 CrossEntropy 会让网络偏向像素多的类别所以工程里常把 CrossEntropy 与 Dice 或 Lovasz 损失线性组合。import torch.nn.functional as F def mixed_loss(logits, target, alpha0.6): ce F.cross_entropy(logits, target) dice soft_dice_loss(F.softmax(logits, dim1), target) return alpha * ce (1 - alpha) * dicealpha 取 0.6 到 0.8 之间比较常见。alpha 太高起不到缓解不平衡的作用太低会让网络在边界上出现锯齿。另外Dice 损失在类别极少时梯度不稳定如果你的训练图片里某个类别只出现在一两张图上训练时不要开 RandomCrop 以外的类别平衡采样。更好的办法是给 CrossEntropy 加 class weightweight 按像素占比取逆平方根。4.4 训练中常见的报错与排查表报错现象可能原因处理方式AssertionError on class number配置 num_classes 和数据集标签最大值不一致检查标签像素值是否从 0 开始连续size mismatch for decoderHRNet 输出通道数和分割头输入不一致查看 dilation 分支的 out_ch 是否等于 head in_chCUDA out of memorybatch_size 或分辨率太大缩小 batch_size、开启 AMP grad scalerloss nan学习率过大或标签有空值降低 lr检查 ann_dir 是否缺失文件精度在验证时特别差训练集和验证集来自不同区域检查数据划分是否按区域而不是按文件随机分一个容易踩的细节是标签灰度值。很多公开遥感标签在图像里把背景设为 255 或 0 的某个固定值与类别 ID 不一致。如果你的源码自带数据加载器可以在 dataset 配置里加一项 reduce_zero_labelTrue让 0 作为 ignore index但要当心它会把背景类从 mIoU 计算里剔除。4.5 用多尺度推理提升小目标召回在验证阶段我会打开多尺度测试。mmsegmentation 里这样设置tta dict( tta_configdict(img_ratios[0.75, 1.0, 1.25], flipTrue), mergemean)多尺度推理不会改变模型参数但能显著提高小目标和边界处的鲁棒性。代价是推理时间和显存增加 3 到 5 倍。如果只是希望在排行榜上好看可以只在验证时开 TTA如果模型要部署到无人机或卫星在轨推理建议不打开 TTA而是把模型导出为静态尺寸的 ONNX。5. 把 AD-HRNet 迁移到自己的数据集滑窗推理与 ONNX 导出迁移到新区域时遥感影像往往远大于训练尺寸。比如原始影像 6000x6000直接送进网络显存会爆。常见做法是滑窗裁剪overlap 至少要覆盖模型感受野的一半通常设置为 128 到 256 像素窗口边缘部分只取中心区域作为有效预测避免拼接痕迹。可以在推理循环里用一个后处理函数把边缘权重衰减到 0。def window_weight(shape, overlap): h, w shape ramp torch.ones(h, w) ramp[:overlap] torch.linspace(0, 1, overlap) ramp[-overlap:] torch.linspace(1, 0, overlap) ramp[:, :overlap] * torch.linspace(0, 1, overlap).view(1, -1) ramp[:, -overlap:] * torch.linspace(1, 0, overlap).view(1, -1) return ramp导出 ONNX 时需要注意两点膨胀卷积在 ONNX 里对应的是 DilatedConvolution注意力里的 AdaptiveAvgPool 在动态输出尺寸下可能不支持。导出时固定输入尺寸为 512x512并使用 dynamic_axes 只 batch 维度可变torch.onnx.export( model, dummy, adhrnet.onnx, opset_version17, input_names[input], output_names[logits], dynamic_axes{input: {0: batch}, logits: {0: batch}})导出后用 onnxruntime 跑一遍比较 logits 的余弦相似度和像素级 argmax 重合率。通常膨胀卷积和 7x7 空间注意力在 CPU/GPU 上都能正确执行但某些部署框架对 dilation 大于 1 的卷积算子实现不完整需要在部署侧检查是否有算子 fallback。如果 TensorRT 里膨胀卷积被替换为 3 次普通卷积推理速度会有明显下降这时可以改用 5x5 普通卷积在小模型上近似再重新蒸馏一次。本文还有配套的精品资源点击获取