基于PyTorch的遥感图像滑坡识别：语义分割、U-Net与实战调优

简介：这是一份基于PyTorch与CNN网络实现的遥感图像滑坡识别项目，面向深度学习课程设计或期末大作业场景。项目不仅包含AlexNet、ResNet等经典卷积网络结构实现，还提供了数据集划分脚本，代码结构清晰，下载后无需修改即可运行，适合快速复现与二次开发。压缩包内共10个文件，分别为7个Python源码、1个.pth训练权重、1个JSON配置及1个.pyc编译文件，整体大小约51.6MB，轻量且易于部署。该项目经由导师指导完成，曾获97分的期末大作业高分评价，具有较高的工程参考价值，目前已有888人学习/下载。读者可借助完整源码与训练好的权重，系统理解遥感图像滑坡识别的建模流程，为课程设计、毕业设计或后续研究提供可落地的基线方案；同时，AlexNet与ResNet的对比实现也有助于掌握卷积网络的典型设计思路。

1. 遥感图像滑坡识别：先把任务定成像素级分割

解压之后先别急着找 train.py。数据目录里通常是一批高分影像和对应的标注文件，标注可能是 shp 矢量，也可能是已经转好的掩膜图片；不管哪种形态，算法眼里它都是一个像素一个像素的标签。遥感图像滑坡识别属于语义分割而不是图像分类：模型对每个像素判断它属不属于滑坡体，输出一张和原图同尺寸的掩膜，而不是一个 0/1 的标量。这一区别决定了后续绝大部分选择——数据怎么裁剪、损失函数用什么、评估指标看什么，全部围绕“逐像素预测”展开。CNN 负责提取纹理、边界和上下文特征，PyTorch 负责把特征计算组织成可训练的图结构并做自动求导。把任务定准了，后面每一步才不走偏。

2. 数据组织与CNN模型选型：U-Net编码器要接得住遥感波段

拿到训练好的模型之前，先要理解数据和模型是怎么配起来的。遥感图像滑坡识别最常见的输入是高分卫星或无人机影像，GeoTIFF 格式，常见 3 到 4 个波段（RGB 加近红外）。滑坡体在近红外波段反射率偏低，多一个波段对模型识别是有实际帮助的，因此很多方案的 Dataset 按 4 波段读取，而不是把它当普通 JPG 直接喂进 CNN。

2.1 影像与掩膜怎么配对：波段、尺寸和文件名规则

一个典型的项目文件夹里，影像和标注通常是这样组织的：

读取时最忌讳直接用 PIL 打开。遥感影像多为 16bit 无符号整型，值域、波段数和地理元数据都和普通图片不一样。常见做法是用 rasterio 读取：

import numpy as np import rasterio def read_image_and_mask(image_path, mask_path): # 以 4 波段为例，读取后转成 float32 便于后续归一化 with rasterio.open(image_path) as src: img = src.read().astype(np.float32) # (C, H, W) with rasterio.open(mask_path) as src: mask = src.read(1).astype(np.uint8) # (H, W) return img, mask

这里src.read()返回的数组维度顺序是(C, H, W)，和 PyTorch 张量期望的维度一致，后面只需要做通道维度检查。掩膜读取时用read(1)只取第一个波段，避免 3 波段标注被当作多通道数据混入训练。滑坡区域的像素值建议在 Dataset 里统一转成 0/1，避免出现 255、1 混用的情况。

归一化这里有个常见误区：不要直接套 ImageNet 的 mean/std。遥感地物的像素分布和自然影像差得很远，尤其近红外波段，ImageNet 统计值基本没有参考意义。我一般会先从训练集随机采样几百张影像，统计每个波段的均值和标准差，把统计结果存成 npy 文件，训练和推理共用同一份。

2.2 为什么用编码器-解码器结构：CNN 在滑坡边界上的作用

滑坡在遥感影像上呈现为边界不规则的裸土区域，颜色和周围植被差异大，但和道路、建筑工地等裸地容易混淆。要区分这些，模型既要看到局部纹理，又要看到大范围的上下文。普通分类 CNN 在最后一个特征图上做全局池化，会丢失像素级位置信息，直接拿来用会导致掩膜边缘糊成一片。编码器-解码器结构则把高分辨率特征逐层保留下来。

常见的选型组合是 U-Net + ResNet 编码器。ResNet 作为 encoder 承担 CNN 卷积神经网络的骨干部分，提取从边缘到语义的多尺度特征；decoder 通过上采样逐级恢复分辨率，并通过 skip connection 把底层细节送到高层。具体到滑坡识别，skip connection 尤其重要，因为滑坡边界本身是连续但破碎的，缩得太狠的小碎块很难在上采样阶段凭空恢复。

如果机器显存有限，第一版可以用 ResNet18 做 encoder，输入尺寸 256×256，batch size 设为 4，一张 12GB 显存的卡基本能跑。数据量不大时，ResNet34 是精度和速度比较平衡的选择。需要注意的一点是，遥感影像的波段数是 4 时，不能直接用默认的 ResNet 权重，因为它第一个卷积层是 3 通道。常见做法是把权重复制到第 4 个通道，或者直接丢掉预训练权重从零训练，显然后者在数据量不足时会更难收敛。

2.3 PyTorch Dataset 实现：裁剪、增强和异步加载

训练时把整幅影像直接塞进网络通常是不现实的。高分影像动辄几千乘几千像素，显存放不下。标准做法是随机裁剪出固定大小的 patch，把影像和掩膜按同样的偏移量裁剪，同时做数据增强。这一步用 albumentations 比手动写变换省心，因为它能保证图像和掩膜使用同一套随机参数。

from torch.utils.data import Dataset import albumentations as A import torch transform = A.Compose([ A.RandomCrop(256, 256), A.HorizontalFlip(p=0.5), A.VerticalFlip(p=0.5), A.RandomRotate90(p=0.5), ]) class LandslideDataset(Dataset): def __init__(self, image_paths, mask_paths, mean, std, crop=256): self.image_paths = image_paths self.mask_paths = mask_paths self.mean = mean self.std = std self.crop = crop def __len__(self): return len(self.image_paths) def __getitem__(self, idx): img, mask = read_image_and_mask( self.image_paths[idx], self.mask_paths[idx]) # 先做通道归一化，再裁剪，避免 albumentations 对 float32 做类型转换时溢出 for c in range(img.shape[0]): img[c] = (img[c] - self.mean[c]) / self.std[c] img = img.transpose(1, 2, 0) # (H, W, C)，albumentations 的输入格式 augmented = transform(image=img, mask=mask) img = augmented['image'].transpose(2, 0, 1) mask = augmented['mask'] return (torch.from_numpy(img.copy()), torch.from_numpy(mask.copy()).float())

代码里把归一化放在 albumentations 之前，是因为 RandomCrop 内部可能对输入做类型检查，float32 图像配合uint8掩膜混用容易出隐蔽的类型错误。mask.copy()是为了避免 torch.from_numpy 拿到只读数组后报错。__getitem__返回前把掩膜转成 float，是因为后续损失函数需要它参与浮点运算。

3. 用PyTorch写训练脚本：损失函数、IoU评估与参数设置

环境层面，建议用 conda 单独建一个环境，Python 3.9 以上，PyTorch 按 GPU 版本安装好之后先跑一次torch.cuda.is_available()确认 CUDA 可用。模型部分不需要自己从零写 U-Net，segmentation-models-pytorch这个库封装了常见 encoder 和 decoder，能把网络定义压缩到几行，切 backbone 也方便。

import torch import torch.nn as nn from segmentation_models_pytorch import Unet model = Unet( encoder_name='resnet34', encoder_weights='imagenet', in_channels=4, # 对应 RGBN 四个波段 classes=1, )

注意这里仍然用了 ImageNet 预训练权重，虽然第二个卷积层只接收前 3 个通道。库在初始化时遇到in_channels=4会自动把第一个卷积层权重复制到第 4 个通道，不用手动处理。这样做的好处是训练前期收敛快，坏处是近红外通道初始特征和真实分布不一致，如果训练数据充足，可以关掉预训练从头训。

3.1 损失函数：BCE 加 Dice 的组合

滑坡掩膜里前景像素占比通常很低，一张 512×512 的影像里滑坡区域可能只占 5%。只使用 BCE 会让模型倾向于把所有像素预测成背景，因为这样损失也很小。常见做法是把 BCEWithLogitsLoss 和 DiceLoss 按 1:1 加权相加，前者提供稳定的梯度，后者直接优化前景和背景的交并比。

class DiceLoss(nn.Module): def __init__(self, smooth=1e-5): super().__init__() self.smooth = smooth def forward(self, logits, targets): prob = torch.sigmoid(logits) # 先展平成 (B, N)，逐样本计算 Dice 再求平均 prob = prob.reshape(prob.size(0), -1) targets = targets.reshape(targets.size(0), -1) intersection = (prob * targets).sum(dim=1) union = prob.sum(dim=1) + targets.sum(dim=1) dice = (2 * intersection + self.smooth) / (union + self.smooth) return 1 - dice.mean()

smooth参数放在分子分母上同时是为了防止除零，也避免训练初期某个样本完全没有前景时梯度剧烈跳动。targets必须是 0/1 的浮点数，不能是 one-hot 后的多通道形式，否则这里维度展平后算出来的 Dice 就失去了逐像素对应关系。

3.2 评估指标：IoU 和 Recision/Recall 怎么算

滑坡识别最终交付时，评估指标不能只看 loss。IoU 是最直观的指标，它把预测掩膜和真实掩膜的交集除以并集。PyTorch 里在验证阶段需要注意关闭梯度，并且用torch.sigmoid(logits) > 0.5得到布尔掩膜再计算。

def compute_iou(pred_mask, true_mask, eps=1e-6): pred_mask = (pred_mask > 0.5).bool() true_mask = true_mask.bool() intersection = (pred_mask & true_mask).sum().float() union = (pred_mask | true_mask).sum().float() iou = (intersection + eps) / (union + eps) return iou.item()

这个实现是纯张量运算，验证时可以直接对整个 batch 做。但它算的是滑坡这一类的 IoU，背景类没有参与。另一种做法是分别计算两个类的 IoU 再取平均，即 mIoU，数据里滑坡占比极低时 mIoU 会比单类 IoU 更保守，也更常用于遥感分割论文里。实际项目中建议两个指标都输出，提交报告时说明清楚。

3.3 训练参数表与启动命令

训练超参不需要太复杂，先按下面这组跑一个 baseline，再根据曲线调整。

训练脚本里，验证函数写在每个 epoch 结束后，只对验证集做推理，不参与反向传播。保存权重时不要只存最后一步，按验证集 IoU 保存最优权重，文件名带上 epoch 和 IoU 数值，比如best_ep37_iou0.823.pth。这样后面加载推理时，能清楚知道手头这个权重是在什么条件下产出的。

4. 加载训练好的模型做推理：state_dict检查与切块预测

拿到训练好的模型文件之后，第一步不是直接跑预测，而是确认权重和当前模型结构对得上。torch.load返回的对象里存的是state_dict，它是一个从权重名到张量的字典。如果之前保存时用了torch.save(model.state_dict(), path)，加载时就要先实例化一个完全相同的 UNet，再用load_state_dict加载。

4.1 加载权重前先检查 state_dict 的键

常见错误是Missing key(s) in state_dict，最常见的原因有两个：一是保存时包了一层model.module（DDP 训练遗留下来的），二是模型定义时类名或参数不一样。建议加载前先打印权重文件的键名看结构。

import torch device = 'cuda' if torch.cuda.is_available() else 'cpu' checkpoint = torch.load('best_ep37_iou0.823.pth', map_location='cpu') if 'state_dict' in checkpoint: state_dict = checkpoint['state_dict'] else: state_dict = checkpoint # 如果键名带 module. 前缀，去掉前缀 if any(k.startswith('module.') for k in state_dict.keys()): state_dict = {k[7:]: v for k, v in state_dict.items()} model = Unet(encoder_name='resnet34', encoder_weights=None, in_channels=4, classes=1) model.load_state_dict(state_dict) model.to(device) model.eval()

map_location='cpu'的作用是先把权重加载到 CPU，再通过model.to(device)搬到目标设备，避免在 CUDA 设备上直接加载时出现显存分配异常。encoder_weights=None表示不需要预训练权重，因为马上要用自己的 state_dict 覆盖。model.eval()必须调用，否则 Dropout 层在推理时仍然生效，同一个输入每次预测结果都不一样。

注意：如果权重是从Unet之外的自定义模型中保存的，千万不要图省事跳过 key 检查。对比 weight 名里的层名，能反推出当初模型的结构，省到的排错时间远大于读日志的时间。

4.2 单张影像推理：预处理与阈值后处理

推理时的预处理要和训练时保持一致：同样的波段顺序、同样的均值方差、同样的输入尺寸。推理代码里最容易踩的坑是输入张量多了一维还是少了一维。PyTorch 期望的输入形状是(B, C, H, W)，单张影像要unsqueeze(0)变成 batch 维度。

import numpy as np import torch import rasterio def predict_single_tile(model, image_path, mean, std, device): with rasterio.open(image_path) as src: img = src.read().astype(np.float32) profile = src.profile for c in range(img.shape[0]): img[c] = (img[c] - mean[c]) / std[c] tensor = torch.from_numpy(img).unsqueeze(0) # (1, C, H, W) tensor = tensor.to(device) with torch.no_grad(): logits = model(tensor) prob = torch.sigmoid(logits).cpu().numpy()[0, 0] # (H, W) mask = (prob > 0.5).astype(np.uint8) # 把预测结果写回 GeoTIFF，保留原影像的地理参考 profile.update(dtype=rasterio.uint8, count=1, compress='lzw') with rasterio.open('pred_mask.tif', 'w', **profile) as dst: dst.write(mask, 1) return mask

torch.no_grad()是推理时省显存的关键，它告诉 PyTorch 不需要记录任何梯度计算图。profile.update保留了原影像的坐标系和分辨率信息，让预测掩膜可以直接叠加到 GIS 软件里查看。阈值 0.5 是常规选择，但如果项目要求宁可多报不可漏报，这个阈值可以往下降，比如 0.3。

4.3 大影像切块预测：避免拼缝伪影

整幅遥感影像通常大于 1024×1024，直接喂给模型不仅显存不够，而且模型在训练时只见过 256×256 的裁剪块，对过大输入的特征响应并不稳定。标准做法是滑窗切块，每块预测完后拼回去。相邻块之间推荐重叠一半，合并时舍弃边缘，只保留每块中心区域的结果。

def predict_full_image(model, image_path, mean, std, device, patch_size=256, stride=128): with rasterio.open(image_path) as src: img = src.read().astype(np.float32) profile = src.profile H, W = img.shape[1], img.shape[2] mask_sum = np.zeros((H, W), dtype=np.float32) count = np.zeros((H, W), dtype=np.float32) with torch.no_grad(): for y in range(0, H - patch_size + 1, stride): for x in range(0, W - patch_size + 1, stride): patch = img[:, y:y+patch_size, x:x+patch_size] for c in range(patch.shape[0]): patch[c] = (patch[c] - mean[c]) / std[c] tensor = torch.from_numpy(patch).unsqueeze(0).to(device) logits = model(tensor) prob = torch.sigmoid(logits).cpu().numpy()[0, 0] mask_sum[y:y+patch_size, x:x+patch_size] += prob count[y:y+patch_size, x:x+patch_size] += 1 prob = mask_sum / np.maximum(count, 1) mask = (prob > 0.5).astype(np.uint8) profile.update(dtype=rasterio.uint8, count=1, compress='lzw') with rasterio.open('pred_full_mask.tif', 'w', **profile) as dst: dst.write(mask, 1) return mask

stride=128对应 patch_size 的一半，重叠区域在mask_sum和count里被多次累加，最终取平均。这样做的好处是每块预测的边界不会因为上下文缺失而产生明显的拼接痕迹。如果图片尺寸不能被 stride 整除，滑窗循环不会覆盖右下角边缘，算法上通常做镜像 padding 或者从底右边界反向再滑一次，保证全图都被预测到。

5. 滑坡识别调优三板斧：样本均衡、后处理与backbone替换

baseline 跑通后，精度卡在某个水平上不去，通常逃不开三个方向：类别不平衡、预测噪声、骨干网络表达力不足。按顺序排查，比盲目调学习率划算得多。

5.1 类别不平衡：从损失函数下手而不是改数据

如果看到训练 loss 不断下降但最终掩膜几乎全黑，说明模型已经学会把所有像素判为背景。Dice loss 对这种极端不平衡比 BCE 稳，但收敛速度稍慢。更激进一点可以换成 Focal Loss，它的调制系数会让模型把注意力集中在难分类的前景像素上，对滑坡这类小目标效果明显。

数据层面，不要把整幅大图全部放进训练集而只靠随机裁剪采样，因为绝大多数裁剪块可能完全落在背景区域。先统计每幅影像掩膜中前景像素占比，把包含滑坡的裁剪块和纯背景裁剪块按 1:3 的比例混合送进训练，效果通常立竿见影。纯背景块不能完全去掉，否则模型会失去对道路、建筑等易混目标的判别能力。

5.2 预测后处理：用连通域过滤碎噪声

滑坡体在空间上是连续的，预测结果里孤立的 5 像素噪声块基本可以判定为误检。用 scipy.ndimage 做连通域分析，统计每个连通域的面积，把面积小于阈值的区域直接置为背景。这个操作不消耗 GPU，纯 CPU 跑得也很快。

from scipy import ndimage def remove_small_objects(mask, min_area=20): # 标记连通域，返回每个区域的编号和个数 labeled, num_features = ndimage.label(mask) cleaned = np.zeros_like(mask) for region_id in range(1, num_features + 1): region_mask = labeled == region_id if region_mask.sum() >= min_area: cleaned[region_mask] = 1 return cleaned

min_area 的取值和影像分辨率强相关。无人机影像地面分辨率 10cm 时，20 像素大概对应 2 平方米；卫星影像分辨率 3m 时，同一个阈值对应 180 平方米。所以这个参数要先换算实际地物面积再看，不能照抄。

5.3 backbone 替换：从 ResNet34 到轻量高效网络

如果数据量足够或者已经做完数据增强，backbone 就是下一个可以优化的大变量。ResNet 系列虽然通用，但在高分辨率遥感场景里参数利用率不算高。近两年的新骨干网络里，CSPNet 一族的思路是把特征图拆成两条路径，一条走梯度流较短的支路，另一条做常规卷积变换，最后拼接。这样做能在保持 CNN 结构不变的前提下减少重复梯度计算，让模型用更少的参数学到同样的表达能力。

这套滑动窗口时序模拟时延，背靠的就是 utf8。

实际切换时，segmentation_models_pytorch的Unet直接支持encoder_name='cspdarknet53'这类新的编码器名字，训练脚本几乎不需要改动。换 backbone 之后建议把学习率调回 1e-4 重新跑一轮，不要用旧权重续训。观察验证集 IoU 和单张推理耗时，如果精度提升不明显但速度慢了，可以退回 ResNet34；如果两者都有提升，说明模型已经吃下了新骨干的红利。滑坡识别这种任务，模型体积往往不是瓶颈，推理速度才是，做了这一步对比之后，整个方案的模型选型才算真正收尾。

本文还有配套的精品资源，点击获取