Pytorch实现FCN语义分割：从数据到训练推理全流程解析

简介：这是一套基于Python与PyTorch实现的FCN语义分割复现项目，面向希望入门语义分割，或将其用于毕设项目、课程设计、工程实训的PyTorch学习者。项目严格按原论文复现了FCN32s、FCN16s、FCN8s与FCNs四种网络结构，并配套完整的PyTorch工程代码：train.py负责训练与可视化，FCN.py中定义了FCN系列网络及VGGNet特征提取骨干，BagData.py实现了Dataset与DataLoader数据读取及变换，onehot.py完成标注图onehot编码。同时附带约80MB的包图片数据集（bag_data与bag_data_mask），可在小样本条件下快速跑通从数据加载、模型训练到结果可视化的完整流程。压缩包共1218个文件，以1201张jpg训练图片为主，另有png标注图、Python源码、运行配置及说明文档，整体大小71.81MB，目录结构清晰，便于读者按模块学习或迁移改造。目前已有215人学习下载，适合需要从零掌握FCN原理并动手实现语义分割的初学者与进阶者。

1. 一个不到 80M 的 FCN 项目，为什么值得拆开看

做语义分割的入门，很多人第一站就是 FCN。但直接拿 VOC 2012 那种完整数据集训练，光是预处理和内存布局就劝退一半人。这个项目用 Pytorch 复现了 FCN32s、FCN16s、FCN8s 和 FCNs，训练数据只有随机背景下的包（bag）图片，整个数据集不到 80M。正是因为它小，你能一眼看完从数据加载、onehot 编码到网络定义、训练循环的完整链路，不像大项目那样到处都是抽象层。对做毕设、课程设计或个人实验来说，这是非常合适的最小复现集。我会按数据处理、模型结构、训练脚本、推理可视化这条线拆，最后给几个实际能用的改进方向。

2. 从 VGG 到 FCN：转置卷积与跳层结构

2.1 先搭一个能输出特征图的 VGGNet

FCN 的一大贡献是把分类网络的最后全连接层替换成卷积层，同时保留 VGG 前五组卷积作为特征提取器。你在 FCN.py 里会看到 VGGNet 这一个类，它本质上是 VGG16 的 1x1 卷积化版本。常见做法是先定义 VGG 各层的通道数和池化位置，再逐层搭建。构造时用一个列表features存放卷积层，再用classifier存放两个 1x1 卷积，输出通道数分别对应类别数。

class VGGNet(nn.Module): def __init__(self, n_classes, pretrained=False, model='vgg16'): super(VGGNet, self).__init__() self.n_classes = n_classes # VGG 配置：每个数字是卷积核数，'M' 表示最大池化 vgg_cfgs = { 'vgg16': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'], } cfg = vgg_cfgs[model] self.features = self._make_layers(cfg) # 1x1 卷积替代全连接，压缩到类别数 self.classifier = nn.Sequential( nn.Conv2d(512, 4096, kernel_size=7), nn.ReLU(inplace=True), nn.Dropout2d(), nn.Conv2d(4096, 4096, kernel_size=1), nn.ReLU(inplace=True), nn.Dropout2d(), nn.Conv2d(4096, n_classes, kernel_size=1), )

代码里值得注意的细节：_make_layers需要把 padding 统一设为 1，以保证 3x3 卷积不改变特征图尺寸；每个卷积后面接 BatchNorm 能提升小数据收敛速度。最后一个classifier的 7x7 卷积对应原始 VGG 里的全连接层，由于输入图像一般被缩放到 224 或 320，这个 7x7 恰好作用在最终 7x7 特征图上。训练时这一层是随机初始化的，而前面的features可以选择加载 ImageNet 预训练权重，迁移学习效果会好很多。

2.2 转置卷积不是反卷积

FCN 上采样用的是转置卷积（transposed convolution），Pytorch 里对应nn.ConvTranspose2d。很多人误以为它是卷积的逆操作，其实它只是在形状上做了"升维"，可学习参数决定了上采样质量。以kernel_size=3, stride=2, padding=1为例，输出尺寸刚好是输入的两倍，计算关系是：

输出边长 = 输入边长 * stride + (kernel_size - 2 * padding) - stride

self.upsample2 = nn.ConvTranspose2d(1024, 512, kernel_size=4, stride=2, padding=1)

这里 kernel_size 为 4、stride 为 2、padding 为 1 是一个常见组合，它能把尺寸翻倍且不产生棋盘伪影。实际训练里转置卷积的初始化很关键，虽然 Pytorch 默认用均匀分布，但更稳的做法是使用双线性插值初始化：把卷积核初始化为双线性滤波器的形状，再在训练中微调。如果从零训练，转置卷积的梯度很容易不稳定，导致特征图出现网格状噪声，这也是为什么小数据更推荐加载预训练 VGG。

2.3 FCN32s、FCN16s、FCN8s 的跳层逻辑

跳层（skip connection）是 FCN 解决粗粒度分割的核心手段。网络从 pool5 得到的特征图是原始图像的 1/32，直接上采样 32 倍，那就是 FCN32s，边界往往很模糊，因为丢失了大量中高层空间细节。改进是融合 pool4 和 pool3 的特征。

• FCN32s：conv7 输出直接上采样 32 倍到原图尺寸。
• FCN16s：conv7 先上采样 2 倍，与 pool4 逐元素相加，再整体上采样 16 倍。
• FCN8s：在 FCN16s 基础上，把相加结果再上采样 2 倍，与 pool3 融合，最后上采样 8 倍。

代码里对应 Pytorch 实现是：

# 假设 self.features 返回了各池化层输出 pool3, pool4, pool5 = features x = self.classifier(pool5) # 1/32 x = self.upsample1(x) # 变为 1/16 x = x + self.score_pool4(pool4) # 融合 pool4 x = self.upsample2(x) # 变为 1/8 x = x + self.score_pool3(pool3) # 融合 pool3 x = self.upsample8(x) # 变为原图 1/1

score_pool4和score_pool3是 1x1 卷积层，把 pool4 和 pool3 的通道数压缩到类别数，这样才能和上采样后的结果相加。融合后每个位置既包含高层语义，又保留浅层边界信息。项目里的 FCNs 你可以理解为对跳层融合的进一步扩展，可以再加一层 pool2，也可以把多个上采样分支并联，论文中并没有统一叫法，我们以自己的代码为准。实际操作中，训练 FCN8s 的收敛速度通常比 FCN32s 快，因为梯度能通过跳层直接回传到浅层。

3. 数据准备：Dataset、DataLoader 与 onehot 编码

3.1 BagData.py 里必须实现的三件事

Pytorch 自定义数据集需要继承torch.utils.data.Dataset，并在__getitem__里同时返回输入图像和掩码。这个项目的 BagData.py 把原图路径和 mask 路径一一对应，比较关键的是对 mask 的处理：掩码不是单通道的 0/1 图，而是多个通道的 onehot 编码。因为一个像素只能属于一个类别，语义分割通常用交叉熵损失，它要求 target 是(H, W)的类别索引图。但代码里用 onehot 编码，意味着训练时可能配合二值交叉熵或者多分类交叉熵的 onehot 输入。

一个典型的__getitem__结构是：

class BagData(Dataset): def __init__(self, img_dir, mask_dir, transform=None): self.img_paths = sorted(glob.glob(img_dir + '/*.jpg')) self.mask_paths = sorted(glob.glob(mask_dir + '/*.png')) self.transform = transform def __len__(self): return len(self.img_paths) def __getitem__(self, idx): image = Image.open(self.img_paths[idx]).convert('RGB') mask = Image.open(self.mask_paths[idx]) # 确保图像和掩码尺寸一致后再变换 if self.transform: image = self.transform(image) mask = self.transform(mask) # mask 是 0/1 单通道，转成 long tensor 即可喂给交叉熵 mask = torch.squeeze(torch.tensor(np.array(mask), dtype=torch.long)) return image, mask

需要注意的一点是：图像用的 transform 里Normalize不能用于 mask。常见的做法是把 mask 的 transform 单独定义，只做 Resize 和 ToTensor。代码里如果 mask 是彩色索引图，ToTensor后变成 3 通道，你要手动取其中一个通道或者用torch.unique检查类别数。这里项目选择了先做 onehot，好处是能呈现更清晰的多通道标签，坏处是会占用额外内存。

3.2 onehot.py：把索引映射成多通道概率分布

def onehot_encode(mask, n_classes): # mask: [H, W], 值域 [0, n_classes-1] onehot = torch.zeros((n_classes, mask.size(0), mask.size(1))) for c in range(n_classes): onehot[c][mask == c] = 1 return onehot

这个实现循环遍历类别，逻辑最简单，但速度慢。更推荐用F.one_hot：

mask_onehot = F.one_hot(mask.to(torch.int64), num_classes=n_classes).permute(2, 0, 1).float()

F.one_hot要求 mask 是整型，输出形状是(H, W, C)，再用permute转成 Pytorch 期望的(C, H, W)。用float()是因为后面可能跟二值交叉熵损失。如果你打算用nn.CrossEntropyLoss，就不需要 onehot，直接把类别索引图传进去就行。这个小项目保留 onehot 编码，很大程度是为了教学演示，让你直观看到 label 长什么样。

3.3 DataLoader 参数设置与验证

data_loader = DataLoader(dataset, batch_size=8, shuffle=True, num_workers=2, pin_memory=True)

num_workers在 Windows 下不要超过 4，否则容易报错；pin_memory=True能减少 GPU 传输时间。如果训练时显存不够，优先降低 batch_size 而不是削减网络层。

4. 训练：train.py 里的超参数、损失函数与循环细节

4.1 一套能跑通小数据的超参配置

这个项目训练集不到 100 张图，网络却是 VGG16 级别的参数。下表是我跑类似项目时的推荐配置：

参数	值	说明
batch_size	4	小数据集太大 batch 容易过拟合
learning_rate	1e-4	初始化用 Adam，后面切 SGD 微调
epochs	100	小数据可以跑久一点，但早停要开
loss	CrossEntropyLoss	如果 mask 是 onehot，则用 BCEWithLogitsLoss
optimizer	Adam	weight_decay 设为 1e-4
lr_scheduler	ReduceLROnPlateau	验证 loss 不降就把 lr 减半
image_size	320x320	太小丢边界信息，太大会撑爆显存

train.py主函数里通常会一次性把网络、损失函数、优化器都建好，然后循环epoch。小数据集的特色是每个 epoch 时间短，你可以打印每个 batch 的 loss，观察抖动趋势。

4.2 训练主循环的写法

for epoch in range(start_epoch, epochs): for i, (images, labels) in enumerate(train_loader): images = images.to(device) labels = labels.to(device) outputs = model(images) # 输出 [N, n_classes, H, W] loss = criterion(outputs, labels) # labels 可以是 [N, H, W] 索引或 [N, C, H, W] onehot optimizer.zero_grad() loss.backward() optimizer.step() if i % 20 == 0: print(f'Epoch {epoch:3d} / {i:4d} loss {loss.item():.4f}') # 每个 epoch 结束做一次验证 val_loss = validate(model, val_loader, criterion, device) scheduler.step(val_loss)

如果labels是 onehot 形式，交叉熵损失不直接支持，你要么把网络输出经过softmax后与 onehot 做BCELoss，要么用torch.argmax(labels, dim=1)转回索引图再用CrossEntropyLoss。项目本身提供 onehot.py，大概率是后者。注意nn.CrossEntropyLoss内部自带 softmax，你把网络末尾不要加 softmax 层，直接输出 logits。

4.3 用 tensorboard 或 matplotlib 做可视化

from torch.utils.tensorboard import SummaryWriter writer = SummaryWriter('runs/bag_fcn') # 在每个 epoch 结束时 writer.add_scalar('Loss/train', train_loss, epoch) writer.add_scalar('Loss/val', val_loss, epoch) for name, param in model.named_parameters(): if param.requires_grad and param.grad is not None: writer.add_histogram(f'grad/{name}', param.grad, epoch)

梯度直方图能帮你发现梯度消失或爆炸。如果大量参数的梯度过早变为 0，说明网络某些分支没有学到东西，常见原因是跳层融合时两个支路数值量级不一致，可以在相加前对其中一个支路做x * 0.1缩放，或者初始化 1x1 卷积的参数。

4.4 过拟合信号怎么识别

小数据集最容易遇到训练 loss 不断下降、验证 loss 却飙升。这时先看训练集里是否真的有几乎相同角度的重复图片。这个 bag 数据集主背景是随机的，所以过拟合主要是记住了特定物体的纹理。你可以增加随机裁剪、翻转和颜色抖动。另一个信号是验证集的 IoU 不升反降，那就把学习率调回 1e-5，再做几次 epoch，看看 loss 有没有触底反弹。

5. 推理与可视化：从 checkpoint 到像素级结果

5.1 加载模型并处理输入

训练保存模型时，建议只保存 state_dict：

torch.save(model.state_dict(), 'fcn8s_bag.pth')

推理时新建一个相同网络，再load_state_dict，注意带上strict=False以防 key 不匹配。输入图像需要做和训练时一样的预处理：Resize 到H x W，转 Tensor，用训练时记录的均值和标准差做 Normalize。

from PIL import Image from torchvision import transforms import torch def preprocess_image(path, size=320): img = Image.open(path).convert('RGB') trans = transforms.Compose([ transforms.Resize((size, size)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) return trans(img).unsqueeze(0)

注意Normalize的均值和标准差要与你训练时的一致，否则输出的 logits 分布会偏离，分割效果急速劣化。如果你训练时没做 Normalize，推理时必须保持一致，不要临时加。

5.2 输出后处理与叠加显示

网络输出[1, n_classes, H, W]，语义分割的预测类别由通道维上的最大值决定，用torch.argmax(outputs, dim=1)得到单通道索引图。这个索引图是 Resize 后的尺寸，要显示到原图上还需要Resize回原始分辨率。

model.eval() with torch.no_grad(): out = model(inputs) # [1, n_classes, 320, 320] pred = torch.argmax(out, dim=1) # [1, 320, 320] pred = pred.squeeze(0).cpu().numpy() # [320, 320] pred_img = Image.fromarray(pred.astype('uint8')) pred_img = pred_img.resize((orig_w, orig_h), Image.NEAREST) # 叠加到原图：将预测区域半透明覆盖 orig = Image.open(path).convert('RGBA') mask = pred_img.convert('L').point(lambda p: 128 if p > 0 else 0) colored = Image.new('RGBA', orig.size, (255, 0, 0)) combined = Image.blend(orig, colored, alpha=0.3)

argmax后的 NumPy 数组用Image.fromarray时要确认 dtype，uint8 可以直接生成灰度图。预测图 Resize 回原图尺寸时要用最近邻插值，避免破坏类别边界。如果想看每个类别的概率分布，比如特定类别是背景还是包，可以输出out[0, class_id]并在 matplotlib 里画热力图。

5.3 计算简单指标：pixel accuracy 和 mIoU

def pixel_accuracy(pred, true): return (pred == true).sum().item() / true.numel() def iou_per_class(pred, true, n_classes): ious = [] for c in range(n_classes): pred_c = (pred == c) true_c = (true == c) intersection = (pred_c & true_c).sum().item() union = (pred_c | true_c).sum().item() if union == 0: ious.append(float('nan')) else: ious.append(intersection / union) return ious

小数据集上 pixel accuracy 容易虚高，因为背景占比大。mIoU 更敏感，如果网络只输出了背景，背景 IoU 可能接近 1，目标类 IoU 接近 0。观察 mIoU 曲线比 loss 曲线更能反映分割质量。

6. 三个立竿见影的改进方向

6.1 用 VGG 预训练权重做初始化

项目代码里的 VGGNet 允许传入pretrained=True，但需要你提前下载 torchvision 提供的模型。加载时只取features部分，因为分类器形状不同：

import torchvision.models as models vgg = models.vgg16(pretrained=True) new_model.features.load_state_dict(vgg.features.state_dict())

这样你的 FCN 一开始就有很好的边缘和纹理特征，训练只需要微调后面几层。小数据集从零训练容易卡在局部最优，预训练权重能明显提升收敛速度和最终 IoU。

6.2 用空洞卷积替换部分池化

VGG 连续五次最大池化让特征图缩小到 1/32，这丢失了大量细粒度信息。一个常见做法是把后两个池化层的 stride 改为 1，同时对后续卷积层使用空洞卷积（dilated convolution）来维持感受野。Pytorch 里只要给nn.Conv2d加dilation参数：

nn.Conv2d(512, 512, kernel_size=3, padding=2, dilation=2)

这样特征图停留在 1/8 或 1/16，最后上采样倍数降低，分割边界更锐利。项目里的 FCNs 如果想升级，可以把这个版本做成"ASPP 风格"的多尺度空洞卷积并联，再融合，效果会比单纯的跳层更好。

6.3 后处理加上 CRF 或测试时增强

要快速提升输出质量，不用改网络也能做两件事：第一，测试时把输入水平翻转、改变尺度各跑一次，平均 logits 后再 argmax；第二，用全连接 CRF 对概率图做区域一致性修正。CRF 实现可以用pydensecrf或opencv-contrib里的ximgproc，它对噪声点清理非常有效，代价是速度慢。实际工程中，如果部署端有 GPU，也可以尝试torchcrf，或者干脆训练一个轻量级残差模块来模拟 CRF 的效果。这个 bag 数据集背景简单，加上翻转和多尺度就已经能看到边界明显规整，推荐先试这个再上 CRF。

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