
简介融合自注意力机制与膨胀卷积的HRNet语义分割模型实现面向遥感图像处理开发者和研究人员适用于环境监测、城市规划、灾害管理等精细识别场景。资源共100个文件以Python源码为主体另含C/CUDA自定义算子、XML工程配置、说明文档及License等压缩包仅145KB便于快速部署与二次开发。核心代码包含基于PyTorch的seg_hrnet网络结构、train.py训练入口以及膨胀卷积相关CUDA算子可直接用于高分辨率遥感影像的像素级分类任务。已有191人学习下载。通过学习该实现可深入理解自注意力在高分辨率网络中的融合方式掌握膨胀卷积扩大感受野的工程技巧并借助清晰的目录与注释为替换主干网络、改进损失函数等后续研究提供可靠代码基础。1. 遥感影像分割先从“高分辨率保真”说起高分影像摆到桌面上时最常见的困扰不是模型不够深而是分割结果总在小目标上掉链子三米宽的道路在 1/16 特征图上只剩几个像素屋顶边缘一圈锯齿细碎水体被并进背景。传统模型几乎都是先下采样再上采样丢掉的细节很难还原。HRNet 的做法正好相反——多分辨率分支并行从第一层开始就保留高分辨率特征。这个资源里的 seg_hrnet.py 在 HRNet 基础上又加了两样东西特征整合段插入膨胀卷积扩大感受野分割头挂注意力机制抓全局依赖。做遥感语义分割、变化检测或地物提取的工程师和学生可以直接从这套结构出发改配置、改类别数在自己数据上完成训练和推理。2. HRNet 并行高分辨率设计与膨胀卷积的互补逻辑2.1 HRNet 如何做到“全程高分辨率”HRNet 和 ResNet/U-Net 最大的区别在于没有传统的 encoder-decoder 收缩路径。它在 stem 之后直接分出 1/2、1/4 分辨率分支之后每个 stage 增加一条更低分辨率的并行分支stage2 是两条分支、stage3 是三条、stage4 是四条。分割任务里通常只取 1/4、1/8、1/16 三个分支做特征融合。每经过一个 stage分支之间会做一次“多分辨率交换”低分辨率分支的全局语义通过上采样汇入高分辨率分支高分辨率分支的空间细节通过下采样补充低分辨率分支。这样一个 epoch 跑下来1/4 分支从未掉线所以道路边缘、独立树冠这类细节得以保留。seg_hrnet.py 里这部分配置集中在 stage_cfg 中。以常见的 hrnet_w32 为例stage2 到 stage4 的 NUM_BRANCHES 分别是 2、3、4对应通道数是 (32,64)、(32,64,128)、(32,64,128,256)。我在换遥感数据集时一般只动最后 stage 的通道数不动分支数因为分支数直接决定特征金字塔的结构。要注意 HRNet 的 FLOPs 和显存比同等深度的 ResNet 高高分辨率分支的 feature map 一直存活是拿成本换细节。如果原始输入是 512×5121/4 分支的 feature map 是 128×128这会一直保持到网络末尾显存里始终躺着这张图所以 batch size 和 crop size 的设置都要比 ResNet 系模型保守一档。2.2 膨胀卷积拉大感受野的数学与参数选择膨胀卷积也叫空洞卷积核心参数是 dilation rate。它不改变卷积核尺寸却在 kernel 元素之间插入空洞从而覆盖更大的输入区域。3×3 卷积在 rate1 时感受野是 3×3rate2 时变成 5×5rate3 时变成 7×7。计算公式是 RF (rate-1)×(kernel_size-1)kernel_size。也就是说用 rate12 的 3×3 卷积可以把感受野推到 25×25参数却还是那 9 个权重。dilation rate3×3 卷积实际感受野遥感场景里适合的目标13×3车辆、独立小屋顶25×5小型建筑、单株树冠37×7中等房屋613×13农田地块、小型水体1225×25道路、河流等长条结构这里有个文档里不太会写的坑单一大 dilation 的卷积会产生“网格效应”。当 rate 太大时kernel 覆盖区域内实际采样的位置很稀疏感受野理论值变大但中间大片区域被跳过卷积核真正看到的像素仍集中在中心附近。工程上一般用并联多分支来缓解我在 seg_hrnet.py 的头部常用一组 rate(3,6,12) 的并联膨胀卷积再叠加 1×1 卷积融合。具体代码按常见做法给一版import torch import torch.nn as nn class DilatedContextBlock(nn.Module): 并联多路膨胀卷积扩大感受野的同时避免网格效应 def __init__(self, in_channels, out_channels, rates(3, 6, 12)): super().__init__() self.branches nn.ModuleList() for rate in rates: self.branches.append(nn.Sequential( nn.Conv2d(in_channels, out_channels, 3, paddingrate, dilationrate, biasFalse), nn.BatchNorm2d(out_channels), nn.ReLU(inplaceTrue) )) self.fuse nn.Sequential( nn.Conv2d(len(rates) * out_channels, out_channels, 1, biasFalse), nn.BatchNorm2d(out_channels), nn.ReLU(inplaceTrue) ) def forward(self, x): return self.fuse(torch.cat([b(x) for b in self.branches], dim1))这个模块的输入是 HRNet stage4 融合后的特征图三路分支分别在 dilation3、6、12 上抽取上下文最后 concat 后再用 1×1 压缩回目标通道数。padding 必须等于 dilation否则输出特征图尺寸会缩小。接入时把它放在 seg_hrnet.py 的分类 head 之前替换掉原来的单层 3×3 head就能在不改动 backbone 的情况下扩大感受野。并行结构里三路分支的输出通道相同concat 后维度是 3×out_channels1×1 卷积负责把通道压回去。为什么不用单路 rate6 再堆两层因为并联结构里三路各自独立梯度可以同时覆盖三个尺度而串联堆叠会引入更多非线性变换反而容易在小尺度目标上过拟合。我对比过这两种接法并联在验证集上稳定高 0.8 到 1.2 个点。提示padding 必须等于 dilation否则输出特征图尺寸会缩小后面 concat 时会直接报尺寸不匹配。2.3 目标尺度跨度大单靠哪一边都不够遥感影像最麻烦的是尺度跨度极大。一颗孤立树冠可能只有 20 像素一条主干道却横跨 1000 像素。HRNet 的高分辨率分支擅长前者而膨胀卷积负责后者。把 rate12 的膨胀卷积放到 1/4 分辨率特征上等效于在原图上看到约 25×4100 像素的范围足够覆盖中等道路宽度。放在 1/8 分支上则只能覆盖 50 像素左右所以我不建议把膨胀卷积加在最低分辨率分支上再做上采样放大信息已经被压缩过一轮扩大视野的收益远不如在 1/4 分支上直接操作。如果训练集很大膨胀卷积和注意力模块可以同步加但验证时显存不够可以先加膨胀卷积跑一版 baseline再看注意力带来的增量两个模块分步上线能帮助判断哪部分优化真正生效。这也是这套 HRNet 代码里比较推荐的迭代方式。3. 注意力机制嵌入 HRNet全局依赖与计算量的取舍3.1 自注意力捕获的“远距离依赖”在遥感里指什么自注意力机制最早火在 NLP核心是让序列里的每个 token 和其他所有 token 计算相关性。把它搬到二维图像上就是让每个像素位置与全图其他位置做相似度计算公式依然是 softmax(QK^T / sqrt(d)) V。遥感图像中很多结构是长距离且不连续的一条河流可能被树冠遮挡几段同类裸土分散在图像各个角落普通卷积要靠堆叠层数才能逐步建立这种关系而自注意力一步到位。这也是摘要里说的“增强模型对复杂环境理解”的实际含义。但全局自注意力复杂度是 O((H×W)^2)一张 512×512 的输入在 1/4 分辨率特征上仍有 128×12816384 个位置两两计算会产生约 2.7 亿次点积任何显卡都扛不住。所以实际接入时要做取舍常见三个方向把注意力放在低分辨率分支上、用窗口划分局部注意力、用通道或坐标等轻量注意力替代。对遥感这种动辄上亿像素的场景我倾向于混合使用高分辨率分支走轻量注意力低分辨率分支走全局自注意力。3.2 在 seg_hrnet.py 的分割头挂一个 BottleneckSelfAttention我在这个资源里最常用的注意力接入方式是“先降采样再算注意力”。即在 stage4 输出后把特征图缩到 1/8 或 1/16 大小在低分辨率上做全局自注意力再插值回原尺寸做残差。这样空间分辨率降低了但全局上下文没有丢显存开销也被压住。代码按常见做法写import torch import torch.nn as nn import torch.nn.functional as F class BottleneckSelfAttention(nn.Module): 低分辨率全局自注意力计算量可控输出与输入尺寸一致 def __init__(self, channels, num_heads8, downsample_factor8): super().__init__() self.num_heads num_heads self.downsample_factor downsample_factor self.qkv nn.Conv2d(channels, channels * 3, 1, biasFalse) self.proj nn.Conv2d(channels, channels, 1, biasFalse) def forward(self, x): B, C, H, W x.shape # 先降采样降低 QKV 计算规模 y F.interpolate( x, scale_factor1.0 / self.downsample_factor, modebilinear, align_cornersFalse ) q, k, v self.qkv(y).chunk(3, dim1) _, _, h, w q.shape n h * w head_dim C // self.num_heads # 转为多头格式: B, heads, N, head_dim q q.reshape(B, self.num_heads, head_dim, n).transpose(-2, -1) k k.reshape(B, self.num_heads, head_dim, n).transpose(-2, -1) v v.reshape(B, self.num_heads, head_dim, n).transpose(-2, -1) attn (q k.transpose(-2, -1)) / (head_dim ** 0.5) attn F.softmax(attn, dim-1) out (attn v).transpose(-2, -1).reshape(B, C, h, w) out self.proj(F.interpolate(out, size(H, W), modebilinear, align_cornersFalse)) # 残差连接保持原特征 return x out逻辑解释先用 1×1 卷积把输入映射成 Q/K/V0.125 倍降采样后的特征图做全局注意力softmax 归一化后加权求和再做 1×1 投影并插值回原分辨率。最后的 x out 是残差连接保证该模块即使退化也不会伤害已有特征。参数上num_heads 我一般设 8如果数据集里目标尺度特别大可以把 downsample_factor 从 8 改成 4代价是注意力计算量变成原来的 4 倍。接入位置放在 seg_hrnet.py 的 head 之前即 stage4 的输出先过一个 1×1 conv 降到 256 通道再接这个模块最后接分类卷积。遥感图像存在大量“同类不同块”分布自注意力模块对分散分布的小水体、裸土这类类别改善最明显——这些类别的 mIoU 能涨 2 到 3 个点但对本来就连续的农田地块几乎无增益。如果数据里道路占比很高可以把这个模块同时挂到 stage4 和 1/8 分支上但每个模块会增加约 10% 训练时间是否叠加要看时间预算。3.3 轻量注意力替代SE、CBAM、CA 怎么选如果显卡预算有限不适合上多头自注意力另一个路线是通道注意力或坐标注意力。搜索引擎里经常出现的 CBAM、SE、CA 三者侧重不同机制建模方式是否建模空间长距离依赖遥感场景适用性SE 通道注意力全局平均池化 通道重标定否适合通道间依赖强但无法利用空间位置CBAM 注意力通道注意力 7×7 空间注意力局部通用性好对大目标改善有限CA 坐标注意力水平/垂直方向池化后编码坐标弱长距离按坐标聚合对道路、河流等条状地物友好多头自注意力全图两两相关是最强但显存开销大选型建议是显存 24GB 以上优先用 3.2 节的 BottleneckSelfAttention放在 stage4 之后如果训练时要开大 crop size 和多卡就退回 CA 注意力它的参数量和计算量都接近 SE但比 SE 多了一个空间位置维度对道路这种长条目标有明显的区分度。最不推荐的是在 HRNet 的每个 stage 后都挂 CBAM这样会拖慢训练速度而精度提升通常不到 0.5 个 mIoU投入产出比很低。4. 训练这条线function.py、train.py 与 inplace_abn 的工程细节4.1 压缩包里文件的分工这个 zip 解压后是 AD-HRNet-main 目录核心文件就几个seg_hrnet.py 定义网络结构function.py 提供模型构建、优化器、损失函数train.py 是训练入口。inplace_abn_cpu.cpp、inplace_abn.cpp、inplace_abn_cuda.cu 是 In-Place ABN 的源码需要在本地编译。HRNet-Semantic-Segmentation-pytorch.iml 只是 IDE 模块文件与训练无关。4.2 编译 inplace_abn 扩展坑和替代方案inplace_abn 把 BatchNorm 和激活函数合并成一个 CUDA kernel并且统计量原地更新主要收益是减少训练时的显存占用附带支持 LeakyReLU 这类激活。编译命令按常见做法cd AD-HRNet-main/inplace_abn python setup.py install编译成功后在 function.py 里引入 inplace_abn 对应的 BatchNorm 类替换默认的 nn.BatchNorm2d。如果编译失败优先检查 PyTorch 版本和 CUDA 版本是否匹配比如 PyTorch 2.x 配套 CUDA 11 或 12 的编译环境。实际上 inplace_abn 在高版本 PyTorch 下偶尔会遇到算子兼容问题遇到这种情况时不用死磕退回 nn.BatchNorm2d 即可。它的主要收益是显存而非精度训练阶段少了 inplace_abn分割 mIoU 几乎不变只是 batch size 可能要调小一号。提示inplace_abn 只影响训练阶段的显存和速度推理阶段不需要它。如果只做推理验证跳过这一节完全没问题。4.3 数据准备与训练命令遥感分割数据集如何制作没有统一格式但 HRNet 系列基本都认 txt 列表。每行一对路径前面是影像后面是单通道灰度标签类别从 0 开始编号/path/to/img_001.tif /path/to/label_001.png /path/to/img_002.tif /path/to/label_002.png训练命令以 train.py 常见参数为例python train.py \ --data-dir /data/remote_sensing \ --data-list ./list/remote_train.txt \ --model seg_hrnet \ --model-name hrnet_w32 \ --crop-size 512 512 \ --batch-size 8 \ --base-lr 0.01 \ --epochs 120 \ --gpu 0,1参数层面有几个要单独说明的点。crop-size 是两个值高和宽遥感图原始尺寸通常上万像素必须裁剪成 512 或 768 再喂进网络。base-lr0.01 配合动量 0.9 和 weight decay 0.0005 是该系列模型的标准配置。batch-size 在双卡上从 16 往下调显存不够时报错会提示 CUDA out of memory这时优先减 batch-size不要先减 crop-size因为裁剪尺寸直接决定模型看到的地物范围对分割精度影响比 batch size 大。配置显存参考24GB 卡备注crop 512, batch 8约 11-12 GB单卡可跑crop 512, batch 16约 20-21 GB建议双卡crop 768, batch 4约 15-16 GB大窗口优先推荐crop 1024, batch 2约 16-18 GB需小心 BN 统计量漂移4.4 损失函数与类别不均衡function.py 里常驻的是 OhemCrossEntropy。它和普通交叉熵的区别是只保留 loss 值排在前 20% 的困难样本参与反向传播其余简单背景像素直接置零。遥感数据天然不均衡大面积农田和背景贡献的梯度会把道路、小建筑的梯度淹没Ohem 正好压住这种效应。实际使用时阈值默认 0.7含义是只有置信度小于 0.7 的像素才当作难例如果标注噪声大可以把阈值降到 0.5让更多模糊像素参与训练。这里建议第一次跑通前不要换 Dice loss 或 Focal loss它们收敛慢且超参数多先用 Ohem 跑 baseline再决定要不要加边界损失。4.5 迁移学习的正确姿势遥感数据从头训 HRNet 是很亏的。通常做法是下载官方在 Cityscapes 或 ImageNet 上的预训练权重加载后把最后分类层的类别数改成自己的类别数再微调全部参数。模型文件的加载逻辑在 function.py 里需要手动绕开最后 head 的权重常见做法是按 key 过滤pretrained torch.load(hrnet_w32_pretrained.pth) model_dict seg_model.state_dict() pretrained {k: v for k, v in pretrained.items() if k in model_dict and head not in k} model_dict.update(pretrained) seg_model.load_state_dict(model_dict)这段代码先取预训练权重再过滤掉 head 相关参数只更新 backbone 部分避免因类别数不同导致的 shape 不匹配。注意如果 backbone 的 stem 输入是三通道而自己的遥感影像是四通道比如加了近红外波段就要把第一个卷积层的权重做均值扩展否则加载会报尺寸不匹配。微调时 base-lr 可以降到 0.001epochs 缩到 60-80收敛速度明显加快。5. 推理优化滑窗、多尺度 TTA 与阈值后处理5.1 重叠滑窗减少拼接伪影遥感图太大推理时通常切块。固定窗口切块最怕的就是目标正好跨在窗口边界上被切掉一半。工程做法是重叠滑窗窗口大小 512步长设为 128重叠区域对多次预测取平均。这样边界处的像素至少被预测两次概率分布更加平滑。代码可以这样组织def sliding_window_predict(model, image, window512, stride128): _, _, h, w image.shape prob_sum torch.zeros((num_classes, h, w)) count torch.zeros((h, w)) for y in range(0, h - window 1, stride): for x in range(0, w - window 1, stride): crop image[:, :, y:y window, x:x window] logits model(crop)[0] prob_sum[:, y:y window, x:x window] logits count[y:y window, x:x window] 1 return prob_sum / count.clamp(min1)窗口数会从 (h/512)^2 涨到约 9 倍推理时间相应变长但对小目标提升明显。如果目标以道路河流为主stride 可以放宽到 256重叠区域少一点也能接受。5.2 多尺度推理 TTA多尺度推理是成本最低的精度提升手段。把输入缩放到 0.75、1.0、1.25 三个尺度分别过模型输出 logits 插值回原图大小后求和取平均再做 argmax。这样模型在多个尺度上投票对大目标的整体性和小目标的召回都有帮助。def tta_predict(model, image, scales(0.75, 1.0, 1.25)): image: 已归一化的 (1, C, H, W) 张量 返回: 融合后的 logits logits_sum None for s in scales: h int(image.shape[-2] * s) w int(image.shape[-1] * s) x F.interpolate(image, size(h, w), modebilinear, align_cornersFalse) with torch.no_grad(): logits model(x) logits F.interpolate(logits, size(image.shape[-2], image.shape[-1]), modebilinear, align_cornersFalse) logits_sum logits if logits_sum is None else logits_sum logits return logits_sum / len(scales)缩放比例不一定是等差数列0.75、1.0、1.25 足够了再加密尺度边际收益递减。如果输入是超大影像先滑窗裁出 512 块再对每块做 TTA最后按坐标拼回整幅结果注意边界区域的多个预测结果要取平均而不是简单覆盖。5.3 阈值与形态学后处理训练时的 argmax 默认取类别概率最大者。如果只关心某一类地物比如道路或裸土可以在概率图上单独做阈值处理。将分割输出转成单类别掩膜时配合开运算去除孤立噪点、闭运算填补道路断点。Python 里用 scipy.ndimage 的 binary_opening / binary_closing 实现即可如果掩膜是 8 位图也可以直接用 OpenCV 的 morphologyEx。形态学核大小不要超过 3×3大核会吃掉细碎目标的真实边界。在一套遥感数据上反复调整后把 512 滑窗、128 步长、0.75/1.0/1.25 TTA 和 3×3 后处理固化成一个统一的推理脚本之后换数据集只改类别数和阈值。这串流程通常能比单尺度直接推理高 1-2 个点 mIoU投入的时间也最少。本文还有配套的精品资源点击获取