企业官网建设流程全解析

简介：面向图像语义分割学习与课程设计的U-Net训练和测试代码压缩包，覆盖数据加载、网络搭建、训练、测试及可视化全部流程。项目共18个文件，以7个Python脚本为核心，分别负责数据预处理、U-Net网络结构定义、训练循环、测试评估与结果展示；同时附带PyTorch权重文件、样例测试图片、XML项目配置以及Markdown说明文档，整体压缩后仅2.15MB，轻量且便于快速部署。代码内注释详细，模块划分清晰，即使刚接触深度学习的新手也能顺着工程骨架理清语义分割的实现思路；对于需要完成期末大作业或课程设计的学生，可直接复用或在此基础上扩展改进。目前已有959人学习下载，该资源的完整性和实用性得到较多验证，是理解U-Net原理、上手PyTorch分割任务的可靠参考。

1. U-Net 图像语义分割落地：这份 PyTorch 代码包能解决什么问题

跑过遥感图像分割的人应该都有同感：U-Net 是少数不需要大规模预训练、在一两百张图上就能出效果的网络结构。但真正动手做期末大作业或者第一个落地版本时，卡住你的往往不是模型结构本身，而是数据路径写错、mask 读成三通道、测试代码没删干净导致训练集混入验证图这类工程问题。这套用 PyTorch 实现的 U-Net 图像语义分割训练与测试代码，就是把这些坑提前填平的完整工程模板。它适合三类人：需要交图像分割期末大作业的学生、刚入门 PyTorch 想找一个能跑通基线的开发者、以及想在遥感语义分割上快速出一版结果的研究生。压缩包里的代码不需要你从零搭，改完路径就能把训练到测试整条链路走通。

2. 训练环境与数据准备：Anaconda 下 PyTorch 选型与 VOC 风格数据集组织

2.1 环境选型：为什么这份代码选择 PyTorch 而不是 TensorFlow

我之前帮人调过几份语义分割代码，最怕的不是 U-Net 写不写得出来，而是训练环境能不能一次配好。这份代码默认在 PyTorch 1.13 + CUDA 11.7 的组合上验证过，用 2.x 的版本也能跑，但千万不要图省事直接pip install torch。那样装下来大概率是 CPU 版，训练一个 epoch 的时间够你泡三杯茶。图像语义分割的训练瓶颈不在模型参数量，而在于框架版本和显卡驱动错位带来的连环报错，尤其是期末赶工时最容易在这里翻车。

conda create -n unet python=3.8 -y conda activate unet pip install torch==1.13.1+cu117 torchvision==0.14.1+cu117 --extra-index-url https://download.pytorch.org/whl/cu117

第一行创建 Python 3.8 的独立环境，避免把系统 Python 搞乱；第二行激活；第三行指定了cu117后缀，意思是这个 PyTorch 版本对应 CUDA 11.7。这里我建议不要用pip install torch不带版本号，因为新版 PyTorch 默认对应更高 CUDA 版本，如果你的显卡驱动是两年内的旧驱动，装了之后torch.cuda.is_available()会返回 False，那一瞬间你会以为显卡坏了，其实是版本错位。

先跑一下nvidia-smi看驱动支持的最高 CUDA 版本，再回头选 torch 的+cu后缀。整个过程的核心原则是：驱动版本决定了你能装什么 CUDA，CUDA 决定了你能装哪个 torch。下面是我常用的版本对照参考：

2.2 数据集目录：VOC 风格怎么组织才能被训练脚本读进去

这份代码用的是 VOC 风格的数据组织方式，也就是把原图和 mask 分开放，再用 txt 文件记录哪些图片属于训练集、哪些属于验证集。我之前见过有人把 mask 直接放到和原图同一个目录里，后缀还全改成.png，最终调试时发现train.txt里写的名字对不上，DataLoader 报了一堆 FileNotFoundError。

datasets/ ├── JPEGImages/ # 原图，jpg 或 png 都行 ├── SegmentationClass/ # mask，必须是单通道灰度图 └── ImageSets/ └── Segmentation/ ├── train.txt └── val.txt

JPEGImages 放原始影像，SegmentationClass 放对应的语义标签。train.txt 和 val.txt 里每行写一个不带后缀的文件名，例如img_001。为什么 mask 必须单通道？因为语义分割的标签本质是一个索引图，每个像素存的是类别编号，不是 RGB 颜色。如果你用标注工具导出了三通道的彩色 mask，需要先转成灰度索引图再放进 SegmentationClass，不然训练时 CrossEntropyLoss 会把每个像素当三个类别算，loss 永远降不到合理区间。

2.3 Dataset 类实现：读图、resize 与 mask 的类型处理

Dataset 类是这份代码里最值得抄的一段，很多报错都藏在容易被忽略的细节里。我重构后的版本长这样：

import os import cv2 import torch from torch.utils.data import Dataset class SegmentationDataset(Dataset): def __init__(self, root, split='train', size=(512, 512)): self.img_dir = os.path.join(root, 'JPEGImages') self.mask_dir = os.path.join(root, 'SegmentationClass') split_file = os.path.join(root, 'ImageSets', 'Segmentation', f'{split}.txt') with open(split_file, 'r') as f: self.names = [line.strip() for line in f.readlines() if line.strip()] self.size = size def __len__(self): return len(self.names) def __getitem__(self, idx): name = self.names[idx] image = cv2.imread(os.path.join(self.img_dir, name + '.jpg')) image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) mask = cv2.imread(os.path.join(self.mask_dir, name + '.png'), cv2.IMREAD_GRAYSCALE) image = cv2.resize(image, self.size, interpolation=cv2.INTER_LINEAR) mask = cv2.resize(mask, self.size, interpolation=cv2.INTER_NEAREST) image = torch.from_numpy(image).permute(2, 0, 1).float() / 255.0 mask = torch.from_numpy(mask).long() return image, mask

代码里两个细节必须说清楚。第一，读 mask 时必须加cv2.IMREAD_GRAYSCALE，否则 cv2 默认读成三通道 BGR，后面所有形状都会乱掉。第二，image 的 resize 用INTER_LINEAR做平滑，而 mask 的 resize 必须用INTER_NEAREST，因为最近邻插值不会产生新像素值，如果用线性插值，mask 边缘会出现 127 这类并不存在的类别编号，训练时 Loss 直接算出一个莫名其妙的数。最后image转成 float 并除以 255，mask转 long，CrossEntropyLoss 的 target 要求的就是 long 类型。

3. U-Net 模型与训练循环实现：从双卷积块到 checkpoint 保存

3.1 U-Net 结构拆解：编码器、解码器与跳跃连接

很多教程喜欢把 U-Net 画成一张大 U 形图，然后告诉你左边是编码器右边是解码器。真正落实到代码，核心就是一个反复调用的双卷积模块。原始论文里只有 Conv + ReLU，但这份代码在每个双卷积块里加了 BatchNorm2d，这个改动很关键，尤其是你的训练集只有几百张图时，BN 能显著抑制网络初始化不稳定带来的 loss 剧烈震荡。

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, kernel_size=3, padding=1), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True), nn.Conv2d(out_ch, out_ch, kernel_size=3, padding=1), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True), ) def forward(self, x): return self.conv(x)

padding=1保证了卷积不改变特征图尺寸，这是跳跃连接能直接相加的前提。如果你把 padding 改成 0，每过一次卷积宽高各减 2，到解码器拼接特征图时维度对不上，会报 size mismatch。BN 层在 batch_size 为 1 的时候会告警甚至报错，所以训练时 batch_size 至少给 2，我一般不低于 4。编码器部分就是连续四次maxpool + DoubleConv，下采样把空间信息压缩成语义信息；解码器部分用转置卷积逐步恢复分辨率，再和对应层级的编码器输出做通道拼接。

3.2 训练循环与超参：CrossEntropyLoss、Adam 与 batch 调参

训练脚本最核心的是循环结构。很多新手喜欢在循环里写一堆 wandb、tensorboard 之类的日志逻辑，导致代码又长又难查。我建议先把最朴素的训练循环跑通，再往上加日志。下面的结构可以直接当模板用：

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = UNet(in_channels=3, num_classes=21).to(device) criterion = nn.CrossEntropyLoss(ignore_index=255) optimizer = torch.optim.Adam(model.parameters(), lr=1e-4) for epoch in range(epochs): model.train() total_loss = 0.0 for images, masks in train_loader: images, masks = images.to(device), masks.to(device) outputs = model(images) # (B, C, H, W) loss = criterion(outputs, masks) # masks: (B, H, W) optimizer.zero_grad() loss.backward() optimizer.step() total_loss += loss.item() print(f"epoch {epoch:03d} loss {total_loss/len(train_loader):.4f}")

重点理解两个地方。一是masks的形状必须是(B, H, W)，不带通道维度。如果 train_loader 返回的是(B, 1, H, W)，criterion 不会报错但会静默算错，因为 PyTorch 在 target 有额外维度时行为很迷。二是num_classes必须包含背景类别，如果你有 20 类目标，输出通道就是 21，mask 中背景像素标 0，目标类别从 1 开始编号。ignore_index=255用来跳过边界标注区域的 loss，VOC 数据集里边界标注通常就是 255。

超参方面我给一个我自己跑遥感分割常用的起点，这个组合在大多数小数据集上不会太离谱：

3.3 checkpoint 保存：只存 state_dict 不存整个 model

保存模型是最容易被忽略的环节，也是期末验收时经常出问题的环节。我见过有人用torch.save(model, 'model.pth')把整个模型对象存下来，然后换了一台机器加载，直接报类定义找不到。正确做法是只保存权重：

best_iou = 0.0 if val_iou > best_iou: best_iou = val_iou torch.save(model.state_dict(), 'checkpoints/best_model.pth')

state_dict是一个 dict，存的是每一层的参数张量。加载时需要先实例化一个结构完全相同的模型，再调用load_state_dict。如果你在训练代码里对模型结构做了小改动，比如加了 Dropout，旧权重文件就加载不上了，会报Missing key(s) in state_dict。所以我会在保存时顺手把模型结构定义也拷贝一份到 checkpoints 目录底下，避免几个月后自己都忘了当初用的什么结构。

4. 训练与测试阶段的常见问题排查：五条血泪踩坑记录

4.1 loss 一直不降，训练了 20 个 epoch 还在 2.3 徘徊

现象：训练一轮下来 loss 从 2.5 缓慢降到 2.3 就不再动了，预测结果几乎全输出同一个类别。

原因：最常见的是学习率太大导致 loss 在小范围内震荡下不去，或者类别严重不平衡，背景像素占了 90% 以上，模型学到的策略是全部预测为背景。另一个隐蔽原因是 mask 的像素值从 1 开始编号而背景是 0，但 CrossEntropyLoss 默认要求类别编号从 0 开始连续分布，如果 mask 里出现了类别编号断层，loss 也会异常。

解决：先看 mask 像素值分布，用np.unique(mask)输出有几类、编号是否连续。然后把 learning_rate 降到 1e-4 重新训练。如果还是不行，检查是不是torch.from_numpy(mask)没有转 long，float 类型的 target 会让 loss 计算出完全错误的值。

4.2 OOM：batch_size 没动过，换了台机器就爆显存

现象：同一份代码，在实验室的 3090 上跑得好好的，换到自己的 2060 上直接RuntimeError: CUDA out of memory，而且不是训练开始时报，是第几个 epoch 之后突然爆。

原因：训练集图片尺寸不统一，Dataset 里 resize 的尺寸是给验证集用的，训练时某个 batch 里的图特别大，显存瞬间吃满。另一个原因是 PyTorch 默认缓存了整个计算图的中间激活值，512×512 输入、U-Net 深层的特征图本身就很大，显存占用远超你按参数量估算的数。

解决：确认 Dataset 的size参数被训练和验证 DataLoader 同时传入了。然后像代码包里那样显式设置torch.cuda.empty_cache()在每个 epoch 结束后调用。最直接的办法是把训练尺寸从 512 降到 384，U-Net 对输入分辨率不敏感，384 和 512 的精度差距很小，显存却能省下将近一半。如果还爆，batch_size 降到 4，不要再降了。

4.3 测试代码没有删干净，训练集里混进了 test 图片

现象：训练时 loss 曲线很漂亮，验证精度也正常，但最终测试结果惨不忍睹，mIoU 比验证低了十几个点。后来一查，train.txt 里混了十几张测试集的图片。

原因：这是最典型的工程事故。很多人写完测试脚本后随手把 test 图片的名字 append 到了 train.txt 里，用来"临时看看训练效果"，然后忘了删。训练阶段模型见过这些图，验证时自然表现好，一到真正的测试集就现原型。我之前排查过一份代码，就是这个原因导致整个训练白跑。

解决：检查ImageSets/Segmentation下的 train.txt 和 val.txt，和测试集文件名做差集，确认没有交集。建议在训练脚本开头加一个断言，当训练集和验证集有交集时直接抛异常：

train_ids = set(open('train.txt').read().split()) val_ids = set(open('val.txt').read().split()) assert len(train_ids & val_ids) == 0, "训练集和验证集存在重叠，请检查 txt 文件"

4.4 加载权重报错：state_dict 的 key 对不上

现象：训练完保存了 best_model.pth，测试时加载报Missing key(s) in state_dict，或者size mismatch for decoder.up4.weight。

原因：保存权重时模型跑在 DataParallel 里，state_dict的 key 会带module.前缀；测试时加载到单卡模型上，key 对不上。另一个原因是训练时输入是 3 通道，测试时改成了单通道灰度图，第一层卷积权重维度就不匹配。

解决：在测试脚本里，如果报module.前缀问题，加载后手动去掉前缀即可。如果报 size mismatch，十有八九是输入通道改了，检查模型的in_channels参数。

state_dict = torch.load('checkpoints/best_model.pth', map_location='cpu') new_state_dict = {k.replace('module.', ''): v for k, v in state_dict.items()} model.load_state_dict(new_state_dict)

4.5 PyTorch 装好后 CUDA 不可用，报 No kernel image available

现象：torch.cuda.is_available()返回 True，但训练第一个 batch 时直接CUDA error: no kernel image is available for execution on the device。

原因：显卡驱动太老，而 PyTorch 对应的 CUDA 版本太高。比如老一点的 GTX 10 系列显卡，装 CUDA 12 的 PyTorch 就会触发这个报错。语义分割训练本身就是显存密集型任务，驱动版本这一层卡住会浪费大量时间。

解决：确认显卡型号后，装一个对应+cu117或+cu118的 PyTorch，这两个版本对老显卡的兼容性最好。装了新版 PyTorch 的话直接降级：

pip install torch==1.13.1+cu117 torchvision==0.14.1+cu117 --extra-index-url https://download.pytorch.org/whl/cu117

5. 测试与评估：单张推理、mIoU 计算与结果可视化

5.1 单张推理：加载权重、预处理、argmax 与颜色映射

测试脚本的逻辑和训练完全不同，没有 backward，也不应该有任何梯度计算。整个推理流程可以拆成四步：加载权重、预处理、前向推理、可视化。下面这段是代码包里 predict 函数的核心，我在它的基础上加了些注释：

def predict_image(model, image_path, device, size=(512, 512)): model.eval() image = cv2.imread(image_path) image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) image = cv2.resize(image, size, interpolation=cv2.INTER_LINEAR) x = torch.from_numpy(image).permute(2, 0, 1).unsqueeze(0).float() / 255.0 x = x.to(device) with torch.no_grad(): logits = model(x) pred = logits.argmax(dim=1).squeeze(0).cpu().numpy() return pred

argmax(dim=1)是在通道维度上取最大值索引，输出形状从(1, C, H, W)变成(1, H, W)，然后再squeeze(0)去掉 batch 维。这里最容易漏掉的是torch.no_grad()，如果不写，推理时 PyTorch 会额外构建计算图，白白吃掉大量显存，跑着跑着就 OOM，而且测试集越大越明显。预测出来的pred是一个二维数组，里面每个像素存的是类别编号，要直接看结果，得做一个颜色映射再存盘。

5.2 mIoU 与像素准确率：用混淆矩阵算，别手动算

判断 U-Net 训得好不好，光看 loss 不够，我一般以 mIoU 为准。mIoU 的计算容易写错的地方在于交集和并集的计算，以及忽略 255 边界像素。这部分代码会直接决定你的验收结果对不对，我建议直接用混淆矩阵实现：

import numpy as np def compute_miou(model, val_loader, num_classes, device): model.eval() confusion = np.zeros((num_classes, num_classes), dtype=np.uint64) with torch.no_grad(): for images, masks in val_loader: images = images.to(device) preds = model(images).argmax(dim=1).cpu().numpy() masks = masks.cpu().numpy() for p, m in zip(preds, masks): valid = (m != 255) p = p[valid] m = m[valid] np.add.at(confusion, m, p) ious = [] for cls in range(num_classes): inter = confusion[cls, cls] union = confusion[cls].sum() + confusion[:, cls].sum() - inter iou = inter / union if union > 0 else 1.0 ious.append(iou) return np.mean(ious)

先说np.add.at这行。它做的事情是把预测类别p作为列索引、真实类别m作为行索引，在混淆矩阵对应位置累加，等价于手动统计每个类别的预测情况。我最早写的是confusion[m, p] += 1，遇到重复索引时不生效，因为 NumPy 的+=不会触发 unbuffered 累加，这个坑非常隐蔽。第二段计算 IoU 时，confusion[cls, cls]是真正预测对的像素数，confusion[cls].sum()是所有真实类别为 cls 的像素数，confusion[:, cls].sum()是所有被预测为 cls 的像素数，相加再减去交集就是并集。没有出现的类别直接给 1.0，避免除零。

6. 进阶技巧：多尺度滑窗推理与 ONNX 模型导出

6.1 多尺度预测：让小目标分割不再"糊"

单尺度推理简单，但遥感分割里小目标占比高，直接 resize 到 512 会把小房子、小汽车抹掉。我一般在这种场景下用多尺度预测，让模型对同一张图按 0.75、1.0、1.5 三个尺度分别推理，再把概率图插值回原尺寸做平均。这个操作不改变训练过程，只改推理阶段，能白捡一到两个点的 mIoU。代价是推理时间变成三倍，批量跑测试时会明显变慢，权衡之后只在精度要求高的场景启用。

6.2 导出 ONNX：把 U-Net 从训练环境里解放出来

代码包里自带了 PyTorch 转 ONNX 的脚本，这是我后来加进去的。模型转成 ONNX 后就不依赖 PyTorch 环境了，C++ 或者服务端可以直接加载推理，这在部署阶段非常实用。导出时唯一要注意的是给 dummy input 一个固定尺寸，U-Net 是下采样四次的结构，输入尺寸必须是 16 的整数倍，否则转出来的模型推理时尺寸稍微不对就直接报错。我的习惯是固定用 512×512 导出，部署端统一在外面做好 resize。

model.eval() dummy = torch.randn(1, 3, 512, 512) torch.onnx.export(model.cpu(), dummy, 'unet.onnx', input_names=['input'], output_names=['output'], opset_version=11)

从第一次在这份代码上调通，到现在我已经养成了一个固定习惯：每次把代码交给别人之前，强制在一个全新的 conda 环境里完整跑一遍训练和测试，顺手把临时测试文件全部删掉，再生成 mIoU 报告。这个流程救了我很多次，有一次就是因为忘了删测试文件，整份代码多花了一周才交出去。这套 U-Net 工程包也是以同样的标准整理好的，希望帮到你。

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