UNet实现遥感图像语义分割：PyTorch毕业设计源码与踩坑实践

简介：一份基于UNet的遥感图像语义分割Python毕业设计项目，含可运行源码与配套论文，面向计算机、地理信息等专业学生，适用于毕业设计、课程设计及期末大作业。项目源码经本地编译运行，评审分达98分，难度适中，覆盖从数据预处理、模型搭建、训练到预测的完整语义分割流程。压缩包共68个文件，约46.94MB，主要包含6个py核心模块、3个ipynb交互式笔记（分别用于数据集制作、模型训练与结果预测）、5个tex论文源文件及对应pdf，另有32个png和6个svg图像辅助说明网络结构及分割效果。论文按毕业设计标准组织，含绪论、相关技术、系统设计、实验分析等章节，可直接参考或扩展。已有369人学习下载，目录提供README、启动脚本与TensorBoard配置，便于快速复现实验、调试参数或在此基础上进行二次开发。

1. 用 UNet 做遥感图像语义分割：这份 Python 毕业设计源码到底能跑出什么

遥感图像和普通照片最大的区别，是目标尺度跨度极大、背景噪声又多，一栋房子可能只占几十个像素，而大片农田又连成一体。如果你正为毕业设计选方向，或者课程设计需要一份能交差的完整工程，基于 UNet 的遥感图像语义分割是性价比很高的选择：结构不复杂、训练收敛快、可视化效果好，答辩 PPT 上放两张对比图就很有说服力。这份 Python 毕业设计源码把整个链路都摊开了——数据裁剪、模型定义、训练、预测、TensorBoard 可视化、LaTeX 论文模板全在一个压缩包里，本地能编译能运行，适合直接照着复现再改成自己的数据集。下面按我拆项目的习惯，从文件结构讲到训练参数，再说清楚哪些地方容易翻车。

2. 先看清压缩包结构再动手：文件清单、启动脚本与论文源码的关系

很多同学拿到 zip 第一件事就是双击 README，这没问题，但 README 往往只写了“怎么跑”，没写“为什么这么组织”。我建议先花十分钟把整个目录过一遍，搞清楚哪个文件负责什么，后面改代码才知道动哪里。

2.1 master 目录与 src 四个模块：data / model / cnn / train 各自该改哪里

解压后是一个 master 工程目录，核心代码在 src 文件夹里，包含utils.py、cnn.py、model.py、data.py、train.py和__init__.py。这个拆分方式很典型，符合 PyTorch 工程的基本约定：

cnn.py和model.py的分工值得留意：cnn.py里是卷积、BN、ReLU 这类基础积木，model.py用这些积木搭出 UNet 的编码器、解码器和跳跃连接。这样拆的好处是，你想把 UNet 换成 Attention UNet 或者 ResUNet 时，只需要在model.py里重新组合cnn.py里的模块，不用重写底层卷积。我刚拿到这个工程时，第一件事就是打开cnn.py看它的conv_block返回的是[feature_map]还是(feature_map, skip)，这决定了跳跃连接怎么接。

2.2 两个 .ipynb 和一个 .ps1：从 create_dataset 到 train 再到 predict 的路径

压缩包里还有demo文件夹，放着create_dataset.ipynb、train.ipynb、predict.ipynb三个 notebook。它们对应三条操作路径：

• create_dataset.ipynb：把原始遥感影像切成小 patch，生成训练用的图片和标签。这是整个项目最容易出错的一步，后面第 3 章详细讲。
• train.ipynb：notebook 版的训练入口。它的代码逻辑和train.py基本一致，区别在于你可以分段执行，方便看中间输出。新手我建议先跑这个，报错时能看到具体在哪一行。
• predict.ipynb：加载训练好的权重，对测试影像做推理，并可视化分割结果。答辩前做演示就靠它。

两个 .ps1 启动脚本是 Windows PowerShell 下用的：start_jupyter.ps1启动 Jupyter Notebook，start_tensorboard.ps1启动 TensorBoard。它们的本质只是把常用命令封装好，比如start_tensorboard.ps1里就是tensorboard --logdir=runs这类操作，免去每次敲一长串参数。如果你在 Linux/macOS 上跑，直接在终端执行同样的命令即可。

2.3 LaTeX 论文源码与 PDF：你拿到的不仅是代码

目录里还有毕业论文文件夹，里面是chap1.tex到chap5.tex、Bibs引用库、Figures配图和字体文件夹，另外还有一份基于 U-Net 网络的遥感图像语义分割_郭子睿.pdf。这说明资源附带的不只是可运行代码，还有一份完整的 LaTeX 论文底稿。毕业设计最头疼的论文格式、图表排版、参考文献格式，这里都有现成模板。我一般会建议先看 PDF 了解论文结构，再对照.tex源码学 LaTeX 排版，不要上来就改内容——先把编译跑通，再一段一段替换成自己的实验数据。你只需要一个 LaTeX 编译环境，比如 TeX Live 或 Overleaf，把整个毕业论文文件夹传上去就能编译出 PDF。

3. 把遥感影像切成模型能吃的样子：create_dataset.ipynb 与数据预处理

遥感影像动不动就是几千乘几千像素，直接整图喂给 UNet，显存立刻爆掉。所以第一步永远是裁剪。create_dataset.ipynb做的就是这件事，但裁剪不是简单切块，里面有几个细节决定训练效果好坏。

3.1 滑动窗口裁剪与重叠：为什么 stride 必须小于 patch_size

假设原图是 4096 x 4096，如果你用 256 x 256 的窗口、步长也是 256，那正好能切成 16 x 16 个 patch，互不重叠。问题在于：一个目标（比如一栋楼）如果恰好落在两个 patch 的边界上，模型训练时只看到它的一半，推理时很容易把另一半认错。常见做法是让窗口重叠，也就是步长小于窗口尺寸。下面的代码演示了等效逻辑：

import os import cv2 import numpy as np from tqdm import tqdm image_dir = 'raw/images' label_dir = 'raw/labels' out_img_dir = 'data/images' out_lab_dir = 'data/labels' patch_size = 256 stride = 224 # 比 patch_size 小 32，即 12.5% 重叠 os.makedirs(out_img_dir, exist_ok=True) os.makedirs(out_lab_dir, exist_ok=True) for name in tqdm(os.listdir(image_dir)): img = cv2.imread(os.path.join(image_dir, name)) lab = cv2.imread(os.path.join(label_dir, name), cv2.IMREAD_GRAYSCALE) h, w = img.shape[:2] idx = 0 for y in range(0, h - patch_size + 1, stride): for x in range(0, w - patch_size + 1, stride): img_patch = img[y:y+patch_size, x:x+patch_size] lab_patch = lab[y:y+patch_size, x:x+patch_size] cv2.imwrite(os.path.join(out_img_dir, f'{name[:-4]}_{idx}.png'), img_patch) cv2.imwrite(os.path.join(out_lab_dir, f'{name[:-4]}_{idx}.png'), lab_patch) idx += 1

逻辑说明：外层循环按 y 方向、内层循环按 x 方向滑动，range(0, h - patch_size + 1, stride)保证 patch 不超出图像边界。重叠部分实际上让同一个目标出现在多个训练样本里，相当于免费的数据增强。参数说明：patch_size=256是 UNet 最常用的输入尺寸，兼顾感受野和显存占用；stride=224意味着 overlap 为 32 像素，对遥感语义分割来说这个重叠比例是够用的。如果你显存紧张，patch_size 可以降到 128，但 stride 最好保持比 patch_size 小一点，比如 112。

3.2 标签掩码的处理：灰度值不等于类别编号

裁剪只是第一步。很多遥感数据集的标签是单通道灰度 PNG，里面每个像素的值代表一个类别编号，比如 0 是背景、1 是水体、2 是建筑、3 是植被。问题在于，你在数据加载时有没有把掩码转成torch.long()类型？PyTorch 的 CrossEntropyLoss 期望输入是[B, H, W]的整数张量，但如果你把掩码读成uint8后直接归一化，模型训练时 loss 会直接报错。更隐蔽的错误是：cv2.imread读出来是 BGR 三通道，如果标签图是三通道但每个通道值相同，看起来是灰度图，实际上 shape 是(H, W, 3)，直接 squeeze 会出问题。常见做法是在data.py里加断言：

assert len(mask.shape) == 2, f'Mask should be 2D, got {mask.shape}'

至于是否需要转 one-hot，取决于你用的是CrossEntropyLoss还是BCEWithLogitsLoss。多分类用前者，整数标签直接喂；二分类如果只是“前景/背景”，标签保持 0 和 1 即可。不需要提前做 one-hot，PyTorch 的CrossEntropyLoss内部会处理。

3.3 数据增强：遥感裁剪后别急着训练

裁剪完的 patch 数量通常很可观，但遥感图像存在明显的类别不平衡——“背景”占了绝大多数像素，“目标”往往只有零星几个。直接训练，loss 会被背景主导。在create_dataset.ipynb里一般会做额外的增强，我在复现时习惯加随机水平翻转、垂直翻转和 90 度旋转，这三种操作不改变语义标签的对应关系，代码也简单：

from torchvision import transforms train_transforms = transforms.Compose([ transforms.RandomHorizontalFlip(p=0.5), transforms.RandomVerticalFlip(p=0.5), transforms.RandomRotation(90), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ])

注意，RandomRotation(90)只旋转 90 的整数倍，不会引入插值误差；如果你用任意角度旋转，标签和图像需要同步做相同变换，还得处理边缘填充问题，容易踩坑。Normalize 用 ImageNet 的均值方差是迁移学习时的常见做法，如果你从零训练，也可以不改。别加RandomResizedCrop，它会改变目标的尺度比例，语义分割任务里容易让模型学到错误的空间关系。

4. 读懂 model.py 里的 UNet 与 train.py 的训练参数：网络结构与超参设计

数据准备好了，接下来是理解模型和训练。很多同学直接python train.py就跑，loss 降不下去也不知道看哪里。这一章把 model.py 里的 UNet 结构和 train.py 里的关键参数讲透，你才能根据 loss 曲线判断是网络问题还是数据问题。

4.1 UNet 的编码器-解码器结构与跳跃连接

UNet 为什么适合遥感分割？因为遥感影像的目标边界往往不清晰，比如水体边界的植被、建筑物阴影里的道路，纯靠卷积下采样会丢失细节。UNet 的编码器逐层下采样提取语义信息，解码器逐层上采样恢复空间分辨率，中间的跳跃连接（skip connection）把编码器每一层的特征图拼接到解码器对应层，相当于给解码器提供了高分辨率的细节线索。

model.py 里的 UNet 实现，核心逻辑大概是这样的简化和说明：

import torch import torch.nn as nn from cnn import DoubleConv, Down, Up class UNet(nn.Module): def __init__(self, in_channels=3, num_classes=4, base_channels=64): super().__init__() self.inc = DoubleConv(in_channels, base_channels) self.down1 = Down(base_channels, base_channels * 2) self.down2 = Down(base_channels * 2, base_channels * 4) self.down3 = Down(base_channels * 4, base_channels * 8) self.down4 = Down(base_channels * 8, base_channels * 8) self.up1 = Up(base_channels * 8, base_channels * 8) self.up2 = Up(base_channels * 16, base_channels * 4) self.up3 = Up(base_channels * 8, base_channels * 2) self.up4 = Up(base_channels * 4, base_channels) self.outc = nn.Conv2d(base_channels, num_classes, kernel_size=1) def forward(self, x): x1 = self.inc(x) x2 = self.down1(x1) x3 = self.down2(x2) x4 = self.down3(x3) x5 = self.down4(x4) x = self.up1(x5) x = torch.cat([x, x4], dim=1) # 注意维度是 channels x = self.up2(x) x = torch.cat([x, x3], dim=1) x = self.up3(x) x = torch.cat([x, x2], dim=1) x = self.up4(x) x = torch.cat([x, x1], dim=1) x = self.outc(x) return x

逻辑说明：编码器四个Down模块逐步把通道数翻倍、空间尺寸减半；最底层的两个base_channels * 8是瓶颈层，没有再做下采样。解码器每个Up模块上采样后，和对应的编码器特征图沿通道维度拼接，所以up2的输入通道是base_channels * 16——来自上层的 8 倍通道和来自跳跃连接的 8 倍通道各占一半。参数说明：base_channels=64是 UNet 论文标准配置，显存吃紧时可以改成 32，但分割精度通常会有可感知的下降；num_classes必须改成你自己数据集的类别数，别忘改这里，忘了就会训练时最后一层输出维度对不上。

4.2 损失函数：CrossEntropy 还是 DiceLoss

train.py 里默认用哪种损失，取决于数据集的类别数量。如果是二分类（背景/目标），常见做法是用BCEWithLogitsLoss；如果是多分类，用CrossEntropyLoss。但遥感分割几乎都会遇到类别不平衡问题，纯交叉熵会让模型倾向于把所有像素都预测成背景。我在实际项目中常用的做法是交叉熵加上 Dice Loss 的加权和：

import torch.nn as nn import torch.nn.functional as F class CombinedLoss(nn.Module): def __init__(self, weight_ce=0.7, weight_dice=0.3): super().__init__() self.weight_ce = weight_ce self.weight_dice = weight_dice def forward(self, logits, targets): # softmax 得到概率 probs = F.softmax(logits, dim=1) # 将标签转为 one-hot，与 logits 形状对齐 targets_onehot = F.one_hot(targets, num_classes=logits.shape[1]).permute(0, 3, 1, 2).float() # Dice 部分：先算每个类别的 Dice，再取平均 intersection = (probs * targets_onehot).sum(dim=(2, 3)) union = probs.sum(dim=(2, 3)) + targets_onehot.sum(dim=(2, 3)) dice = (2.0 * intersection + 1e-6) / (union + 1e-6) dice_loss = 1.0 - dice.mean() # 交叉熵部分 ce_loss = F.cross_entropy(logits, targets) return self.weight_ce * ce_loss + self.weight_dice * dice_loss

逻辑说明：Dice Loss 直接优化目标区域的重叠度，对前景占比小的类别比交叉熵更鲁棒。代码里intersection / union逐类别计算，+ 1e-6是为了防止除零。参数说明：weight_ce=0.7, weight_dice=0.3是我常用的比例，如果你发现模型预测结果中目标区域太小，可以把 dice 权重调到 0.5 甚至 0.6。注意，F.one_hot要求标签是整数且在[0, num_classes)范围内，命令行训练时如果标签里有像素值为 255 的 ignore 区域，得先把它掩掉或替换成 0。

4.3 训练超参：学习率、batch size 与模型保存

train.py 里的超参是影响收敛的关键。我复现这个项目时，第一次跑就发现 loss 在 0.69 附近震荡不下降，原因是学习率设成了1e-2，在 Adam 优化器下直接发散。后来改回1e-4才正常。稳定配置参考如下：

python train.py --epochs 50 --batch_size 8 --lr 1e-4 --val_ratio 0.2

参数说明：--epochs 50对遥感分割差不多够用，如果数据量大或者类别多，可以放到 100，但要注意保存最优模型，否则后 30 个 epoch 可能过拟合。--batch_size 8是 4GB 显存的下限，如果你用 24GB 显存的显卡可以调到 16，收敛速度会快不少。--lr 1e-4是 Adam 的典型安全值，配合ReduceLROnPlateau每 5 个 epoch 损失不降就乘以 0.5，实战中效果好于 cosine 退火。--val_ratio 0.2表示从训练集里留 20% 做验证，不要在训练完抽样验证，容易把随机性当精度。

train.py 里模型保存逻辑一般长这样：

best_loss = float('inf') for epoch in range(epochs): train_loss = run_epoch(model, train_loader, optimizer, criterion) val_loss = evaluate(model, val_loader, criterion) if val_loss < best_loss: best_loss = val_loss torch.save(model.state_dict(), 'best_model.pth') print(f'epoch {epoch}: save best model, val_loss {val_loss:.4f}')

这里有个值得注意的细节：保存的是model.state_dict()而不是整个模型对象。state_dict 只包含参数，体积小，加载时只需要UNet(...).load_state_dict(torch.load('best_model.pth'))。如果你保存整个模型对象，每次 PyTorch 版本升级都要重新 pickle，跨环境迁移会非常痛苦。

5. 避坑：遥感语义分割训练里最容易翻车的五个问题

这个项目我在本地反复跑了好几轮，也帮人排查过各种问题。下面按“现象 → 原因 → 解决”列出最值得注意的五条。

5.1 训练时 loss 一直是 0.69 左右，几乎不下降

现象：模型训练十几个 epoch，loss 稳定在 0.69 附近，无论学习率怎么调都纹丝不动。原因：0.69 是二分类交叉熵的初始概率，说明模型把所有像素都预测成了占主导的类别，多半是标签文件没正确加载，掩码全黑，模型学到了“输出全 0”。解决：用cv2.imread读一张标签图，打印np.unique(mask)，确认类别值分布。如果全为 0，去create_dataset.ipynb检查标签路径是否写对；如果标签值范围不对，在data.py里做一次mask = (mask == 255).astype(np.int64)之类的映射。

5.2 输入输出尺寸对不上，forward 报错

现象：torch.cat时两个张量空间尺寸不一致，报错信息在model.py的torch.cat([x, x4], dim=1)这一行。原因：UNet 编码器每层下采样 2 倍，如果输入尺寸是奇数，比如 255 x 255，下采样后变成 127 或 128，跳跃连接两侧尺寸总是差 1 像素。解决：把输入统一处理成 2 的幂次，比如 256、512、1024。在data.py的数据集类里，transforms.Resize((256, 256))是最省事的办法。也可以给 UNet 加 padding，但改到一半很容易引入更大麻烦，不建议。

5.3 显存不够，OOM 报错

现象：CUDA out of memory，训练刚启动就退出。原因：patch_size 太大、batch_size 太高、原图通道数比预期多。解决：先按--batch_size 2跑通，再逐步翻倍；如果 patch_size 是 512，降到 256。还有一种情况是图像 PNG 其实带 4 通道 RGBA，3 通道的 UNet 输入会报Expected 3 channels, got 4，读图时指定cv2.IMREAD_COLOR强制转成 BGR。

5.4 本地能跑，服务器上 import 直接报错

现象：代码在 Windows 本地正常，传到 Linux 服务器python train.py报ModuleNotFoundError: No module named 'cnn'。原因：项目放在 src 目录，直接python train.py时 Python 不会把当前目录自动加入 sys.path，跨平台时相对导入失效。解决：在 train.py 顶部加sys.path.append(os.path.dirname(os.path.abspath(__file__)))，或者始终从 src 目录下用python -m train启动，让 Python 把当前目录当模块根。

5.5 TensorBoard 没有曲线输出

现象：start_tensorboard.ps1启动后浏览器打开端口，一片空白，提示No dashboards are active。原因：训练时没有写日志。PyTorch 不像 Keras 自动记录，必须显式创建SummaryWriter并把标量写进去。解决：在train.py里加上：

from torch.utils.tensorboard import SummaryWriter writer = SummaryWriter('runs/exp1') # 每个 epoch 结束后 writer.add_scalar('loss/train', train_loss, epoch) writer.add_scalar('loss/val', val_loss, epoch)

确认runs目录里有文件后再启动 TensorBoard，不要手滑把日志写到了别的路径。写完后随手writer.close()释放资源，不然多次训练会把日志文件弄乱。

6. 验证与复现：从 TensorBoard 曲线到 IoU 指标，确认你的模型真的训起来了

模型训练完、best_model.pth 也保存了，下一步就是确认它能用——这一步很多人会跳过，直接拿去交差，结果答辩时发现模型在验证集上效果惨不忍睹。

6.1 先看 TensorBoard 的损失曲线

训练完第一时间打开 TensorBoard，不看最终数值，看曲线形态。正常的曲线是训练集和验证集的 loss 同步下降，最后收敛在一个平台期；如果验证集 loss 在某个 epoch 后反弹上升，说明过拟合了，回到--epochs减少训练轮数或加 dropout。如果两条曲线一直贴得很近，且最终 loss 不高，这是一个好信号。

6.2 用 predict.ipynb 做一次完整的推理可视化

预测和训练最大的区别是推理时不需要标签，而且通常不需要裁剪成 patch——你可以整图喂进去，也可以切块再拼回来。predict.ipynb 里一般会演示单张图像的加载和可视化分割结果：

from model import UNet import torch import matplotlib.pyplot as plt model = UNet(in_channels=3, num_classes=4) model.load_state_dict(torch.load('best_model.pth', map_location='cpu')) model.eval() # 加载测试图并转为 tensor img = cv2.imread('test.png') img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img_tensor = torch.from_numpy(img_rgb.transpose(2, 0, 1)).float().unsqueeze(0) / 255.0 with torch.no_grad(): logits = model(img_tensor) # [1, 4, H, W] pred = torch.argmax(logits, dim=1) # [1, H, W]，每个像素是类别编号 pred_np = pred.squeeze(0).numpy() plt.imshow(pred_np, cmap='tab20') plt.savefig('prediction.png', dpi=300, bbox_inches='tight')

在这段代码里，map_location='cpu'是为了在没有 GPU 的机器上也能复现，这是我从一次答辩现场学来的教训，项目复制到教室电脑上跑不起来就尴尬了。torch.argmax在通道维度取最大值下标，得到的是类别标签图。保存成 PNG 时最好用cmap='tab20'，否则类别多时颜色相近容易看不出边界。

6.3 算一次 IoU，别只看 loss

loss 下降不代表分割效果好，尤其类别不平衡时，loss 降了但目标区域全是错的。常见做法是算每个类别的 IoU 和 mIoU：

def compute_iou(pred, target, num_classes=4): ious = [] for cls in range(num_classes): p = (pred == cls) t = (target == cls) inter = (p & t).sum() union = (p | t).sum() if union > 0: ious.append(inter / union) return sum(ious) / len(ious)

灰度分割图的 IoU 大于 0.5，说明模型具备基本识别能力；如果某个类别 IoU 低于 0.2，多半是样本太少或者标签有误，回去查数据而不是调模型。从那以后，我每次复现这类项目都会强制走一遍推理、IoU、可视化三步，确认数据流从头到尾没有断裂，再谈调参。这套流程不复杂，但能帮你避开大多数“训练一时爽，答辩火葬场”的尴尬，希望帮到你。

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