简介:这份资源面向遥感图像处理、深度学习图像分割方向的学生与研究人员,提供基于注意力增强卷积的ResUNet道路提取与语义分割完整实现,适合作为相关专业毕业设计或课程设计的参考方案。压缩包共9个文件,以7个Python源码为主,涵盖模型定义、注意力增强卷积模块、数据增强生成器、IoU损失函数与评价指标等核心环节,另附txt依赖清单与md说明文档,整体约13KB,结构紧凑便于快速理解项目脉络。资源已在本地部署验证可运行,界面美观、操作简单,实用价值较高。目前已有71人学习下载。读者可借助源码与部署教程复现注意力机制与残差结构结合的分割流程,掌握遥感道路提取中特征聚焦与复杂背景鲁棒性的实现思路,并在此基础上开展二次开发或对比实验,为城市规划、交通管理、灾害监测等应用场景提供技术参考。
1. 遥感道路提取为什么总在树荫和立交桥下翻车
做遥感图像道路提取的同学,大概率都经历过这种场景:模型在开阔路段跑得漂漂亮亮,IoU 刷到 0.8 以上,一切换到树荫遮挡、立交桥多层交叉、窄巷子这些地方,道路直接断成几截,或者把停车场、水泥屋顶误判成路。这不是数据不够,也不是训练轮数不够,而是普通卷积的感受野和上下文建模能力,在细长、连续、易遮挡的道路目标上天然吃亏。
ResUNet 本身是个很稳的骨架——残差连接解决深层退化,U-Net 的编码解码结构保留多尺度信息,在遥感语义分割里一直是性价比很高的选择。但它的瓶颈也很明显:卷积核是局部的,判断一个像素是不是道路,靠的是周围一小圈特征,遇到遮挡就抓瞎。注意力增强卷积(Attention-augmented Convolution)要解决的正是这件事——让网络在卷积的同时,学会「往哪里看」,把远处路段的信息拉过来辅助当前像素判断。
这篇笔记围绕《基于注意力增强卷积的 ResUNet 用于遥感图像道路提取和语义分割》这个方向,把本地部署、数据准备、模型结构、训练参数、推理可视化和踩坑排查讲透。适合两类人:一是做毕业设计、需要一套能跑通、界面能看的完整方案的同学;二是想把注意力机制真正落到分割任务里、不想只调包的工程师。读完你应该能自己在本地把整套流程复现出来,并且知道每个参数动了会发生什么。
2. 注意力增强卷积到底增强了什么:结构拆解与选型理由
2.1 从普通卷积到注意力增强卷积,差在哪一步
普通卷积的操作可以写成:对输入特征图每个位置,用固定权重核做加权求和。权重是训练出来的,但对所有位置共享。这意味着网络对「哪里重要」没有显式建模能力,只能靠堆深度隐式学习。
注意力增强卷积的思路是:在卷积输出之后,接一个轻量的注意力分支,对特征图做通道和空间两个维度的重标定,再把注意力权重乘回主分支。常见实现有两种路线:
- 通道注意力(类似 SE 模块):对每个通道算一个 0~1 的权重,决定哪些通道更重要。
- 空间注意力:对每个像素位置算权重,决定哪些区域更值得关注。
在道路提取任务里,空间注意力更关键,因为道路是空间上连续的结构,遮挡区域需要靠远处上下文补。实际落地时,我一般用「通道 + 空间」串联的混合注意力,放在 ResUNet 编码器的每个残差块之后,解码器上采样前再放一次。这样既控制参数量,又能让浅层特征保留边缘细节、深层特征聚焦道路主体。
选型理由很直接:纯 Transformer 分割方案精度可能更高,但显存和训练成本对毕业设计级别的机器不友好;纯 ResUNet 又不够。注意力增强卷积是在两者之间取平衡,改动小、可解释、容易调。
2.2 ResUNet 主干与注意力模块的拼接位置
ResUNet 的结构分三部分:编码器(下采样 + 残差块)、瓶颈层、解码器(上采样 + 跳跃连接)。注意力模块插在哪里,直接影响效果和训练稳定性。我的经验是:
| 插入位置 | 作用 | 参数量增幅 | 推荐度 |
|---|---|---|---|
| 编码器每个残差块后 | 强化多尺度特征筛选 | 中 | 高 |
| 瓶颈层后 | 全局上下文建模 | 低 | 高 |
| 解码器上采样前 | 恢复空间细节时抑制噪声 | 中 | 中 |
| 跳跃连接上 | 过滤编码器传来的冗余信息 | 低 | 中 |
实际部署时,编码器 + 瓶颈层是必加的,解码器视显存决定。跳跃连接上加注意力容易导致梯度不稳,新手先别碰。
2.3 用 PyTorch 搭一个可运行的注意力增强残差块
下面是一个可以直接抄的注意力增强残差块实现,包含通道注意力和空间注意力两个分支:
import torch import torch.nn as nn import torch.nn.functional as F class ChannelAttention(nn.Module): """通道注意力:对每个通道学习一个权重""" def __init__(self, channels, reduction=16): super().__init__() # 全局平均池化 + 两层全连接,压缩再恢复通道维度 self.avg_pool = nn.AdaptiveAvgPool2d(1) self.fc = nn.Sequential( nn.Linear(channels, channels // reduction, bias=False), nn.ReLU(inplace=True), nn.Linear(channels // reduction, channels, bias=False), nn.Sigmoid() ) def forward(self, x): b, c, _, _ = x.size() y = self.avg_pool(x).view(b, c) # [B, C] y = self.fc(y).view(b, c, 1, 1) # [B, C, 1, 1] return x * y # 广播乘回原特征 class SpatialAttention(nn.Module): """空间注意力:对每个像素位置学习一个权重""" def __init__(self, kernel_size=7): super().__init__() # 沿通道维做平均和最大池化,拼成 2 通道再卷积 self.conv = nn.Conv2d(2, 1, kernel_size, padding=kernel_size // 2, bias=False) self.sigmoid = nn.Sigmoid() def forward(self, x): avg_out = torch.mean(x, dim=1, keepdim=True) # [B,1,H,W] max_out, _ = torch.max(x, dim=1, keepdim=True) # [B,1,H,W] y = torch.cat([avg_out, max_out], dim=1) # [B,2,H,W] y = self.sigmoid(self.conv(y)) # [B,1,H,W] return x * y class AttentionResBlock(nn.Module): """注意力增强残差块:卷积主分支 + 通道/空间注意力""" def __init__(self, in_ch, out_ch, stride=1): super().__init__() self.conv1 = nn.Conv2d(in_ch, out_ch, 3, stride, 1, bias=False) self.bn1 = nn.BatchNorm2d(out_ch) self.conv2 = nn.Conv2d(out_ch, out_ch, 3, 1, 1, bias=False) self.bn2 = nn.BatchNorm2d(out_ch) self.ca = ChannelAttention(out_ch) self.sa = SpatialAttention() self.relu = nn.ReLU(inplace=True) # 捷径分支:通道或尺寸不匹配时用 1x1 卷积对齐 self.shortcut = nn.Sequential() if stride != 1 or in_ch != out_ch: self.shortcut = nn.Sequential( nn.Conv2d(in_ch, out_ch, 1, stride, bias=False), nn.BatchNorm2d(out_ch) ) def forward(self, x): residual = self.shortcut(x) out = self.relu(self.bn1(self.conv1(x))) out = self.bn2(self.conv2(out)) out = self.ca(out) # 通道重标定 out = self.sa(out) # 空间重标定 out = out + residual # 残差相加 return self.relu(out)逻辑说明:主分支走两次 3x3 卷积,注意力分支先做通道加权再做空间加权,最后和捷径分支相加。参数说明:reduction=16控制通道注意力的压缩比,通道数少时可改成 8;kernel_size=7是空间注意力的卷积核,太大增加计算量,太小感受野不够,7 是常用折中。stride控制下采样,编码器每经过一个 stage 设一次 2。
提示:注意力模块放在 BatchNorm 之后、ReLU 之前效果更稳,因为此时特征分布已经归一化,注意力权重不会因为数值范围差异被带偏。
3. 本地部署全流程:从环境到推理界面跑通
3.1 环境依赖与版本选择
本地部署最容易翻车的地方不是模型,是环境。CUDA、PyTorch、cuDNN 三者版本必须对齐,否则要么装不上,要么训练时莫名其妙报错。我一般按这个组合来:
| 组件 | 推荐版本 | 说明 |
|---|---|---|
| Python | 3.8 ~ 3.10 | 3.11 以上部分库轮子不全 |
| PyTorch | 1.12 ~ 2.0 | 带 CUDA 版本,别装 CPU 版 |
| CUDA | 11.3 / 11.6 / 11.8 | 跟 PyTorch 官方对应表走 |
| OpenCV | 4.5+ | 读遥感图、做可视化 |
| NumPy | 1.21 ~ 1.24 | 2.x 和部分老库不兼容 |
安装命令按 CUDA 版本选,比如 CUDA 11.8:
# 创建独立环境,避免污染系统 Python conda create -n roadseg python=3.9 -y conda activate roadseg # 安装带 CUDA 11.8 的 PyTorch pip install torch==2.0.1 torchvision==0.15.2 --index-url https://download.pytorch.org/whl/cu118 # 其余依赖 pip install opencv-python numpy tqdm matplotlib pillow gradio逻辑说明:先建独立环境,再装 PyTorch,最后装业务依赖。参数说明:--index-url指定 PyTorch 官方 CUDA 轮子源,别用默认源,否则可能装成 CPU 版。装完用torch.cuda.is_available()验证,返回 False 就是版本没对齐。
3.2 数据集组织与预处理
遥感道路提取常用的是公开航拍数据集,格式一般是「原图 + 二值掩膜」。目录按下面组织,训练脚本直接按文件夹读:
dataset/ ├── train/ │ ├── images/ # 原图 .tif 或 .png │ └── masks/ # 掩膜,道路为 255,背景为 0 ├── val/ │ ├── images/ │ └── masks/ └── test/ ├── images/ └── masks/预处理做三件事:尺寸统一、归一化、数据增强。遥感图尺寸大,直接缩放会丢细节,常用做法是切 patch,比如 512x512,重叠 64 像素。
import cv2 import numpy as np import os def slice_image(img_path, mask_path, save_dir, patch=512, stride=448): """把大图切成重叠 patch,stride 小于 patch 保证重叠""" img = cv2.imread(img_path, cv2.IMREAD_COLOR) mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) h, w = img.shape[:2] idx = 0 for y in range(0, h - patch + 1, stride): for x in range(0, w - patch + 1, stride): img_patch = img[y:y+patch, x:x+patch] mask_patch = mask[y:y+patch, x:x+patch] # 只保留道路占比大于 1% 的 patch,减少负样本 if mask_patch.mean() / 255 < 0.01: continue cv2.imwrite(os.path.join(save_dir, 'images', f'{idx}.png'), img_patch) cv2.imwrite(os.path.join(save_dir, 'masks', f'{idx}.png'), mask_patch) idx += 1 return idx逻辑说明:双重循环按 stride 滑动切图,重叠区域保证边缘道路不被切断。参数说明:patch=512是训练输入尺寸,显存小就降到 256;stride=448表示重叠 64 像素,重叠越多样本越多但冗余也大;0.01是负样本过滤阈值,道路稀疏的数据集可以调到 0.005。
3.3 训练脚本与关键参数
训练主循环要关注四件事:损失函数、优化器、学习率调度、验证指标。道路提取是二分类且正负样本极不平衡,交叉熵会被背景淹没,我一般用 Dice + BCE 组合损失。
import torch import torch.nn as nn class DiceBCELoss(nn.Module): """Dice 损失 + BCE 损失,缓解正负样本不平衡""" def __init__(self, weight_bce=0.5): super().__init__() self.bce = nn.BCEWithLogitsLoss() self.weight_bce = weight_bce def forward(self, pred, target): # pred 是 logits,target 是 0/1 掩膜 bce_loss = self.bce(pred, target) pred_sig = torch.sigmoid(pred) # 展平后算 Dice pred_flat = pred_sig.view(-1) target_flat = target.view(-1) inter = (pred_flat * target_flat).sum() dice_loss = 1 - (2 * inter + 1e-6) / (pred_flat.sum() + target_flat.sum() + 1e-6) return self.weight_bce * bce_loss + (1 - self.weight_bce) * dice_loss逻辑说明:BCE 提供稳定梯度,Dice 直接优化重叠度,两者加权。参数说明:weight_bce=0.5是平衡系数,道路细且少时可降到 0.3,让 Dice 主导;1e-6防止除零。
训练超参我一般这样设:batch_size=8(512 patch,8G 显存够用)、lr=1e-3(AdamW)、weight_decay=1e-4、epochs=100、余弦退火调度。验证指标用 IoU 和 F1,别只看 loss,loss 降但 IoU 不涨说明过拟合了。
3.4 推理与可视化界面
部署教程里说的「界面美观、操作简单」,落地一般用 Gradio 或 Streamlit 做本地 Web 界面。Gradio 更轻,几行就能起一个上传图片、返回分割结果的页面。
import gradio as gr import torch import cv2 import numpy as np # 加载训练好的模型 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = build_resunet_attention(num_classes=1).to(device) model.load_state_dict(torch.load('best_model.pth', map_location=device)) model.eval() def predict(input_img): """输入 RGB 图,返回叠加了道路掩膜的可视化图""" img = cv2.resize(input_img, (512, 512)) img_norm = img.astype(np.float32) / 255.0 tensor = torch.from_numpy(img_norm).permute(2, 0, 1).unsqueeze(0).to(device) with torch.no_grad(): logits = model(tensor) prob = torch.sigmoid(logits)[0, 0].cpu().numpy() mask = (prob > 0.5).astype(np.uint8) * 255 # 红色叠加显示道路 overlay = img.copy() overlay[mask > 0] = [255, 0, 0] vis = cv2.addWeighted(img, 0.6, overlay, 0.4, 0) return vis demo = gr.Interface(fn=predict, inputs=gr.Image(), outputs=gr.Image(), title='遥感道路提取') demo.launch(server_name='0.0.0.0', server_port=7860)逻辑说明:读图、归一化、转 tensor、前向推理、阈值二值化、叠加可视化。参数说明:0.5是分割阈值,召回不够就降到 0.4,误检多就升到 0.6;server_port=7860是 Gradio 默认端口,本地浏览器访问即可。
4. 避坑与排查:那些让训练白跑的细节
4.1 掩膜值不是 0/1 导致 loss 爆炸
现象:训练第一个 epoch loss 就是 NaN 或者几万。原因:掩膜图读进来是 0/255,没归一化就喂给 BCEWithLogitsLoss,target 超出 [0,1] 范围。解决:读掩膜后统一mask = (mask > 127).astype(np.float32),确保只有 0 和 1。
4.2 数据增强把道路几何结构破坏了
现象:训练集 IoU 很高,验证集一塌糊涂。原因:用了随机旋转 90 度以外的角度、随机裁剪太狠,把细长道路切断,模型学到的是碎片。解决:遥感道路增强只用水平翻转、垂直翻转、90 度整数倍旋转,别用任意角度旋转和强透视变换。
4.3 注意力模块加太多导致显存溢出
现象:训练到一半 CUDA out of memory。原因:每个残差块都加通道 + 空间注意力,中间特征图通道数大,注意力分支的额外计算和缓存把显存吃满。解决:只在编码器后两层和瓶颈层加注意力,浅层高分辨率特征图不加;或者把batch_size从 8 降到 4,配合梯度累积。
4.4 学习率太大让注意力权重震荡
现象:loss 曲线剧烈抖动,IoU 上不去。原因:注意力模块的 Sigmoid 输出对输入敏感,初始学习率 1e-2 会让权重在 0 和 1 之间反复跳。解决:初始学习率设 1e-3 或更低,加 warmup,前 5 个 epoch 从 1e-5 线性升到 1e-3。
4.5 推理时忘了 model.eval() 和 no_grad
现象:单张推理结果和训练时验证结果差很多,或者显存越跑越少。原因:没切 eval 模式,BatchNorm 用当前 batch 统计量;没加 no_grad,计算图一直累积。解决:推理前固定model.eval(),推理包在with torch.no_grad():里。
5. 把 IoU 再往上推两个点的几个实操技巧
模型跑通之后,真正拉开差距的是后处理和训练策略的细节。这里说几个我自己反复验证过、成本低但收益明显的做法。
第一个是多尺度推理融合。单尺度推理对宽窄不一的道路适应性差,把同一张图缩放到 0.75、1.0、1.25 三个尺度分别推理,再把概率图缩回原尺寸取平均,再二值化。代价是推理时间翻三倍,但 IoU 通常能涨 1~2 个点,尤其是窄路和边缘区域。代码上就是在 predict 函数里套一层循环,把prob累加后除以尺度数。
第二个是形态学后处理。模型输出难免有零星噪点和断口,用开运算去噪、闭运算补断。关键是结构元尺寸要跟道路宽度匹配,太大反而把窄路抹掉。我一般用 3x3 的核做一次开运算,再用 5x5 做一次闭运算,参数在验证集上试出来。
import cv2 import numpy as np def postprocess(mask, open_k=3, close_k=5): """形态学后处理:先开运算去噪,再闭运算补断""" kernel_open = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (open_k, open_k)) kernel_close = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (close_k, close_k)) mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel_open) mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel_close) return mask参数说明:open_k控制去噪强度,噪点多就加大到 5;close_k控制补断能力,道路断口长就加大到 7,但别超过道路宽度的两倍。
第三个是难例挖掘。训练几个 epoch 后,把验证集里 IoU 最低的样本挑出来,看是树荫遮挡还是立交桥,然后针对性补充这类场景的 patch 进训练集。这比盲目加数据有效得多。我一般每 20 个 epoch 做一次难例分析,把 IoU 低于 0.5 的样本单独存一个文件夹,人工看一遍再决定补什么。
第四个是损失函数里的边界加权。道路边缘像素最容易错,可以在 BCE 里给边缘像素更高权重。做法是对掩膜做一次形态学梯度得到边缘带,生成权重图,边缘处权重设 3~5,其余为 1。这个改动小,但对边缘 IoU 提升明显。
最后说个血泪教训:别一上来就堆注意力模块和复杂损失,先把基线 ResUNet 跑通、IoU 稳定在合理区间,再逐个加改进,每次只动一个变量,记录验证集指标。我见过太多人一次性改五个地方,结果指标掉了都不知道是哪个模块的锅。训练日志、配置文件、随机种子都存好,这是你唯一的后悔药。希望帮到你。
本文还有配套的精品资源,点击获取