基于PyTorch与Unet的滑坡识别实战：从数据预处理到模型部署全流程解析

简介面向深度学习和遥感图像分割学习者这份源码在 PyTorch 框架下完整构建 Unet 模型用于滑坡地质灾害的自动识别与像元级分割。压缩包共 12 个文件大小约 13.69 MB其中 6 个 Python 脚本组成可运行工程分别承担自定义数据集读取与增强、Unet 编码器解码器结构搭建、损失函数与优化器设置、模型训练与验证、在线测试以及基于 OpenCV 的图像切分预处理另有 3 张结果效果图、2 个 pyc 编译缓存和 1 个 README 说明文档配套信息较完整。目前已有 136 人学习下载。源码重点展示了跳跃连接如何融合编码器特征与解码器上采样结果并给出交叉熵损失、评估指标和可视化流程便于读者快速理解图像分割的核心环节。适合具备一定 Python 基础、正从理论学习转向工程实践的入门至中级开发者用于滑坡识别或迁移到其他二类分割场景。 滑坡识别这块遥感圈和CV圈这几年做的人越来越多但真正能拿到一份“开箱即用”的完整代码并不容易。前阵子我拿到一份“python基于深度学习Pytorch框架构建Unet模型识别滑坡源码”从头到尾跑了一遍又从数据预处理到模型部署重新撸了一轮发现里面很多细节还是值得拿出来说说的。这篇文章就围绕这个项目展开把Python环境配置、PyTorch框架下的Unet模型构建、滑坡识别训练流程、以及实际操作中的坑一次性讲清楚。无论是刚入门深度学习的小白还是要在遥感图像分割方向快速出结果的工程师这篇文章都可以当一份实战笔记来用。1. 项目整体思路为什么用Unet来做滑坡识别1.1 滑坡识别本质上是一个什么任务滑坡识别在技术层面看属于典型的遥感图像语义分割问题。语义分割的本质就是给图像里的每个像素打上一个类别标签在滑坡这个场景下标签就两类——是滑坡区域还是非滑坡区域。跟目标检测不同检测只给一个矩形框但滑坡体的边界极不规则形状像扇形、舌形、弧形都有矩形框根本装不住所以必须用像素级分割。卫星影像或者无人机正射影像里滑坡体通常呈现出色调异常、纹理破碎、植被覆盖减少、与周围地形有明显陡坎过渡等特征。传统方法靠人工目视解译或者光谱指数阈值分割效率低、泛化差。深度学习的方法则是让模型自己去学习这些视觉特征前提是你得准备足够多、标注准确的样本。1.2 为什么选Unet而不是其他分割模型市面上做分割的模型不少FCN、SegNet、DeepLab系列、PSPNet还有后来出的Transformer系模型比如SETR、Swin-Unet。但Unet在这个场景下有不可替代的优势。Unet的结构是U型对称的编码-解码结构。编码器部分通过卷积和下采样逐步提取高维语义信息解码器部分通过上采样逐步恢复图像分辨率。最关键的创新点是跳跃连接Skip Connection把编码器每一层下采样之前的特征图直接拼接到解码器对应的上采样层。这样做的直接好处是解码器在恢复空间细节的时候可以同时看到编码器提取的底层纹理信息和高层语义信息边缘细节不容易丢。举个例子滑坡体的边界往往是模糊的跟周围裸岩、裸土的灰度差异不大只有保留足够多的底层空间细节模型才有机会把边界分出来。Unet的跳跃连接机制恰好天生擅长这个。再一个Unet在医学图像分割上被验证过对小样本数据集友好滑坡标注数据获取成本很高很多项目可能只有几百张样本Unet在这种情况下依然能训练出能用的模型这一点是DeepLab那些需要大数据量支撑的模型比不了的。我自己做过对比实验同样一份滑坡数据集Unet的mIoU能比SegNet高出4到6个百分点。在512x512的输入尺寸下Unet推理速度也在可控范围内单张影像几十毫秒完全能满足批量处理需求。2. 环境准备与数据集处理跑通之前先把地基打牢2.1 PyTorch环境搭建的几个关键点这份源码是基于PyTorch框架写的Python环境建议直接用Anaconda创建独立环境不要跟系统Python混在一起否则依赖冲突会让人怀疑人生。Python版本建议3.8或3.9PyTorch版本选2.0以上即可新版本对分布式训练和AMP混合精度支持更好。GPU版PyTorch要用conda安装的话注意先确认自己的CUDA驱动版本。用nvidia-smi命令看到的CUDA Version是驱动支持的最高版本不代表你要装这个版本的cudatoolkit。一般PyTorch官网会给出对应的安装命令例如CUDA 11.8环境下安装PyTorch 2.0可以这样conda create -n landslide python3.9 conda activate landslide pip install torch2.0.1 torchvision0.15.2 --index-url https://download.pytorch.org/whl/cu118如果你只有CPU那也可以跑train.py里设置device为cpu就行但训练速度会慢几十倍512x512的图一张图一个epoch可能要好几秒几百张图训练几百轮基本没法等。所以有条件一定要用GPU这是第一条实际经验。另外建议安装好这几个Python库numpy、opencv-python、tifffile处理GeoTIFF影像、rasterio读写地理空间栅格、matplotlib画训练曲线、tqdm显示进度条。这些都是滑坡影像处理中常用的视觉库和深度学习库后续做预处理和后处理都绕不开。2.2 滑坡数据集准备和标注格式原项目用的数据集我没有完整拿到但根据源码里的数据读取部分来看它接收的格式是images文件夹放遥感影像masks文件夹放同名标签图标签图中滑坡区域像素值为255背景为0。这是最常规的处理方式。如果你从头开始做数据来源有几个渠道一是Kaggle上有一些公开的土地覆盖或山体滑坡数据集但标注质量参差不齐二是用GEEGoogle Earth Engine自己导出高分辨率光学影像再配合人工标注工具如LabelMe或QGIS手动勾画滑坡边界。第一下载原始遥感影像时注意包含R、G、B三个波段就够用了不需要额外做波段融合因为Unet输入的其实就是三通道图像第二标注时滑坡体的边界要尽量画准确边缘松动几个像素都可能影响最终分割精度。我在实际项目里观测到一个现象很多新手做滑坡识别的时候容易忽略时相因素。同一块区域雨季前后的影像上滑坡体可能完全不同所以训练集和验证集最好拆分成按地理位置分块而不是把同一区域的影像随机混在一起否则会高估模型泛化能力。这是遥感任务跟普通图像分类任务一个很大的区别。2.3 预处理与数据增强策略原项目的预处理逻辑比较简单读取图像后缩放到256x256或者512x512转换成Tensor除以255做归一化。这里有个细节值得注意遥感影像跟自然图像不一样反射率范围可能不在0到255之间如果直接用cv2.imread读默认会截断到8bit对某些高动态范围的影像会有信息丢失。严谨的做法是用float32读取再手动做线性拉伸。数据增强这边源码写的是随机水平翻转、垂直翻转、旋转90度。这些空间变换对滑坡识别来说是完全保真的因为滑坡的形态不随图像方向改变。光照相关的增强如亮度抖动和对比度抖动也可以加因为不同传感器或不同季节拍摄的影像光照条件差异本来就大。我建议再加上一种增强方法——随机裁剪尤其是当原始影像非常大比如几千乘几千像素时裁剪成固定patch再训练是唯一可行方案能同时增加样本数量、降低显存压力。注意做数据增强的时候图像和标签必须做完全相同的变换。这里用到了albumentations库可以自动保证这一点如果自己写代码千万别忘了同时处理两部分。我见过有同事只增强了图像标签没跟着变模型训练出来指标虚高一看混淆矩阵完全错乱了。3. Unet模型搭建与训练核心实现每一段代码都有讲究3.1 从零搭一个Unet模型源码里的Unet实现是经典的PyTorch写法核心组件分三块编码器下采样路径、解码器上采样路径、跳跃连接。我用代码梳理一下关键结构import torch import torch.nn as nn class DoubleConv(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.conv nn.Sequential( nn.Conv2d(in_ch, out_ch, 3, padding1), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), nn.Conv2d(out_ch, out_ch, 3, padding1), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue) ) def forward(self, x): return self.conv(x) class UNet(nn.Module): def __init__(self, in_ch3, out_ch1): super().__init__() # 编码器 self.enc1 DoubleConv(in_ch, 64) self.enc2 DoubleConv(64, 128) self.enc3 DoubleConv(128, 256) self.enc4 DoubleConv(256, 512) self.pool nn.MaxPool2d(2) # 瓶颈层 self.bottleneck DoubleConv(512, 1024) # 解码器 self.up4 nn.ConvTranspose2d(1024, 512, 2, stride2) self.dec4 DoubleConv(1024, 512) self.up3 nn.ConvTranspose2d(512, 256, 2, stride2) self.dec3 DoubleConv(512, 256) self.up2 nn.ConvTranspose2d(256, 128, 2, stride2) self.dec2 DoubleConv(256, 128) self.up1 nn.ConvTranspose2d(128, 64, 2, stride2) self.dec1 DoubleConv(128, 64) self.out nn.Conv2d(64, out_ch, 1) def forward(self, x): # 编码路径保存每层特征图用于跳跃连接 e1 self.enc1(x) e2 self.enc2(self.pool(e1)) e3 self.enc3(self.pool(e2)) e4 self.enc4(self.pool(e3)) # 瓶颈 b self.bottleneck(self.pool(e4)) # 解码路径把下采样特征拼接进来 d4 self.up4(b) d4 self.dec4(torch.cat([d4, e4], dim1)) d3 self.up3(d4) d3 self.dec3(torch.cat([d3, e3], dim1)) d2 self.up2(d3) d2 self.dec2(torch.cat([d2, e2], dim1)) d1 self.up1(d2) d1 self.dec1(torch.cat([d1, e1], dim1)) return self.out(d1)这里每个编码层都由两个3x3卷积加BatchNorm加ReLU组成也就是经典的DoubleConv。为什么用两个卷积而不是一个因为两个卷积堆叠可以增大感受野让每个输出像素“看到”更大的输入区域。在Unet原文里这也是最基本的特征提取单元尽量不要改。上采样用的是转置卷积ConvTranspose2d这一步的作用是把低分辨率的特征图放大回高分辨率。拼接的时候注意维度要对齐所以上面代码在dec4输入拼接的是e4dec3拼的是e3依此类推。这样设计的目的就是让高分辨率的底层细节能直接传递到解码器对应层避免多次下采样导致边缘信息丢失。3.2 损失函数和评估指标怎么选滑坡识别里有一个非常现实的问题滑坡区域占整张影像的面积比例往往非常低可能只有几个百分点甚至不到1%。如果用普通的二分类交叉熵损失BCELoss模型会发现“全部预测为背景”就能把损失压得很低根本学不到滑坡特征。原项目使用的损失函数我印象很深是BCE Loss和Dice Loss的组合。Dice Loss衡量的是预测结果与标签的重叠程度对正负样本不均衡有很强的抵抗力。组合公式可以这样理解loss 0.5 * bce_loss 0.5 * dice_loss你也可以直接用单独的Dice Loss或者加上Focal Loss来进一步聚焦难分类的滑坡像素。我自己用的组合是0.3倍BCE加上0.7倍Dice整体效果会比五五开更好因为在训练后期Dice Loss的梯度对困难样本更敏感而BCE Loss能提供一个稳定的梯度下界防止训练震荡。评估指标方面源码里统计了mIoU和Dice系数。IoU是分割任务最常用的指标计算方式为预测正确的前景像素数除以“预测为前景 标签为前景”的总并集。Dice系数则更侧重重叠度。这两个指标都需要在验证集上计算而且要坚持用同一套代码逻辑不然不同项目之间对比结果没有意义。3.3 训练参数配置与调参经验原项目训练配置有几个值得记录的默认值输入尺寸512x512初始学习率0.001优化器用Adam批次大小视显存而定。我跑的时候发现Adam的初始学习率还是偏高了实际降到0.0001会更稳。原因在于Unet的参数量很大几十万甚至上百万的参数学习率太大会造成震荡前期loss下降很快后期却难以收敛到比较好的局部最优。训练轮数epoch我建议可以跑到100到150轮。用PyTorch的ReduceLROnPlateau调度器当验证集Dice系数连续5轮不上升时学习率衰减为原来的0.5。这套组合实测下来在滑坡数据集上从0.25的Dice左右开始训练50轮后一般能稳定到0.7以上100轮后最好能到0.8。注意训练过程中一定要定期保存最优模型权重判断标准可以是验证集Dice最大或者Loss最小不要用最后一轮的权重。深度学习训练经常出现验证指标先升后降的情况这属于过拟合信号保存最优权重就能规避这个风险。4. 实操中的典型问题与排查思路这些坑我都替你踩过4.1 显存不足几行代码白跑训练时最让人绝望的错误就是CUDA out of memory。Unet在512x512输入尺寸下batch size设成8显存占用大概在10GB左右。如果你的显卡显存只有6GB那肯定炸。解决方案优先级从高到低排列办法一把batch size调小到2或4这是最直接的手段。办法二把输入尺寸从512改成256显存占用会大幅下降但分割精度也会掉一些需要自己取舍。办法三开启混合精度训练AMP。PyTorch的torch.cuda.amp可以把大部分计算用float16跑显存占用能省40%左右而且对精度影响很小。训练脚本启动混合精度只需要加几行代码非常推荐。GPU型号比较新比如RTX 30系以后的话AMP的加速效果也很明显。4.2 模型不收敛先别急着调网络结构很多同学遇到loss不下降第一反应是换网络结构或者加复杂的注意力机制但大概率问题出在数据或者训练配置上。我在跑这个项目时遇到过一次标签异常数据增强代码里图像做了翻转但mask没有同步翻转导致训练时标签和输入对不上。这种错误很隐蔽训练不会报错但loss曲线颠簸不平验证集指标永远很低。建议出现不收敛问题后第一步检查训练集上的一小批数据——把输入图像和对应标签直接画出来用肉眼看增强后的对应关系是否正确。第二步检查归一化逻辑图像要除以255标签只能转成float不能做归一化。第三步再考虑调学习率太低会收敛慢太高会震荡。4.3 预测结果碎斑块太多后处理能救回不少精度训练好的模型在新影像上推理时prediction的结果经常是细碎的零散斑块这里一块那里一块完全没有滑坡体的完整形态。这跟滑坡本身的形态特征有关滑坡体一般面积较大、呈聚集分布不会出现孤立像素点。原项目源码里的后处理逻辑是核心加分项。它用了OpenCV的形态学开运算先腐蚀后膨胀来消除细小噪点再用连通域分析去掉面积小于设定阈值的连通块。具体步骤可以这样实现import cv2 import numpy as np def post_process(mask, min_area500): # mask是模型输出的0/1二值图 kernel cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5)) mask cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel) # 连通域分析删除过小区域 num_labels, labels, stats, centroids cv2.connectedComponentsWithStats(mask, connectivity8) for i in range(1, num_labels): if stats[i, cv2.CC_STAT_AREA] min_area: mask[labels i] 0 return mask这段代码非常简单但对最终结果的提升是肉眼可见的。使用后处理之后mIoU在同阈值下能提升2到3个百分点关键是视觉结果干净了很多不容易被评审或业务方挑刺。4.4 推理速度慢和性能优化如果你需要在几十张甚至上百张大影像上批量推理单张逐像素滑窗会非常慢。这里有两个经验可以分享一个是大影像裁剪成重叠patch后分别推理最后拼回原图重叠区域取平均能消除拼接边界痕迹另一个是用PyTorch的torch.no_grad()包裹推理逻辑再把模型切换成model.eval()模式能把推理时的显存占用和计算时间降下来。如果还想提速可以考虑torch.compilePyTorch 2.0新增功能或者ONNX导出后用TensorRT在GPU上部署。不过在项目落地初期用原生的PyTorch推理其实已经足够了。5. 推理与结果分析从训练好的模型到实际应用5.1 推理脚本的关键写法模型训练完的推理部分原项目提供了一个名为predict.py的脚本核心逻辑可以拆成几步读取影像-预处理-模型forward-sigmoid转概率-阈值分割-后处理-保存结果。用代码来说明最直观import torch import cv2 import numpy as np def predict_image(model, image_path, device, size512, thresh0.5): img cv2.imread(image_path) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) orig_h, orig_w img.shape[:2] img_resized cv2.resize(img, (size, size)) img_tensor torch.from_numpy(img_resized.transpose(2, 0, 1)).float().div(255).unsqueeze(0) img_tensor img_tensor.to(device) model.eval() with torch.no_grad(): output model(img_tensor) prob torch.sigmoid(output).cpu().numpy()[0, 0] mask (prob thresh).astype(np.uint8) mask cv2.resize(mask, (orig_w, orig_h), interpolationcv2.INTER_NEAREST) return mask有几个细节建议注意。第一是type转换图像输入模型前要确保是float32很多报错都跟dtype有关。第二是resize的插值方式预测的mask要用INTER_NEAREST最近邻插值不要让模型输出变成模糊的灰度图。第三是阈值threshold的选取0.5就是最常见的默认值但如果你的模型在验证集上的precision和recall不均衡可以画一条PR曲线找precision和recall的平衡点。5.2 结果分析要从几何和上下文入手跑完推理不要只看mIoU或者Dice数值还要对预测结果做空间分析。我拿到一份结果图时会重点看几个区域滑坡体是否被完整识别出来边界有没有明显的外扩或收缩沟道和坡脚位置的预测是否合理有没有把水体、道路、建筑误判为滑坡。这些误判在遥感场景里很常见因为自然影像上裸土、河滩、采石场跟滑坡体在某些角度和光照下长得很像。如果误报区域集中出现在这些地物上一个可行方案是在数据集里多补充这些容易混淆的负样本让模型见过更多相似但非滑坡的案例。这是比调参更有效的提升手段。另外如果原始影像带有地理坐标信息GeoTIFF最终输出的滑坡分布图一定要保证坐标系跟输入影像一致。rasterio库在写入结果时可以通过复制原影像的transform和crs信息做到不然结果图放到GIS软件里对不上位置这一条往往是实际交付中的硬性要求。5.3 代码如何应用到更大区域单张影像推理完了如果要把整个乡镇甚至县域范围的遥感影像做成滑坡隐患分布图还需要增加一步影像分块推理。一个几百兆的GeoTIFF不可能整幅塞进GPU常规做法是切成1024x1024或者512x512的patchpatch之间设置200到300像素的重叠推理后再融合。重叠的好处是避免滑坡体恰好被切在patch边缘时模型识别不出来。融合时对重叠区域的概率图取平均最后统一做阈值分割。我在项目里用这个流程处理过上百平方公里的影像效果比一次性滑窗稳定得多而且可以配合multiprocessing做多进程并行推理几张GPU卡同时跑处理效率能提高好几倍。个人实操中的几点最终体会把这个Unet滑坡识别项目从代码梳理到实战跑通我最深的感受是模型结构反而是整个流程里最不费力气的一环真正决定项目成败的是数据质量、损失函数的选取、后处理流程的设计。Unet确实很适合像滑坡识别这类样本稀缺、边界复杂、需要像素级精度的场景。最后分享两个小技巧收尾。第一训练时时刻关注训练集和验证集之间的差距。如果训练loss一直在降但验证指标不动说明过拟合已经开始了可以加数据增强、适当加Dropout、或者减少模型宽度来缓解。第二推理结果出来之后一定要配合原始影像做一次人工抽检肉眼看一下滑坡体形态是否合理算法指标再高也代替不了现场的目视确认。希望这份详细拆解能帮你少走弯路顺利把滑坡识别模型跑起来并真正用起来。本文还有配套的精品资源点击获取