作者:SkyXZ
CSDN:SkyXZ~-CSDN博客
博客园:SkyXZ - 博客园
一、什么是FCN?
FCN即全卷积网络(Fully Convolutional Networks) ,由Jonathan Long、Evan Shelhamer和Trevor Darrell于2015年在CVPR会议上发表的论文《Fully Convolutional Networks for Semantic Segmentation》中首次提出,是深度学习首次真正意义上用于语义分割任务 的端到端方法。FCN的提出具有里程碑意义,奠定了现代语义分割网络架构的基础。传统的卷积神经网络(如AlexNet、VGG、GoogLeNet)大多用于图像分类任务,其最终输出是一个固定维度的类别概率向量 ,对应整张图像的分类结果。然而,语义分割任务需要为图像中的每一个像素 赋予语义标签,也就是像素级的密集预测(dense prediction),这与分类任务在输出形式上存在根本不同。为了让CNN能够完成这一类任务,FCN在结构上提出了两个关键的创新:
- 将分类网络全连接层改为卷积层 :传统CNN中的全连接层要求输入为固定尺寸(如224×224),这是因为全连接层将所有空间信息扁平化,破坏了图像的空间结构。而FCN的做法是:将全连接层视为特定感受野大小的1×1卷积操作。换句话说:分类网络原来的
fc6
和fc7
层可以被替换为两个大核的卷积层而输出仍然是特征图,但保持了空间维度,只是经过多次下采样后尺寸较小同时模型不再限制输入图像尺寸 ,可以接受任意大小的输入图像,得到相应尺寸的输出特征图(称为score map或heatmap)。这种转换后的网络就称为"Fully Convolutional"------完全由卷积、池化、激活等空间保留操作构成,没有任何破坏空间结构的层(如全连接、flatten等)。 - 使用上采样还原图像分辨率 :由于卷积神经网络中通常包含多个步长大于1的池化操作(如最大池化),在逐层处理过程中图像的空间分辨率会不断降低。例如,在经典的VGG16结构中,最终输出的特征图相较于输入图像被下采样了32倍(例如输入为512×512,则输出仅为16×16)。为了将这类低分辨率的类别图还原为与原图相同大小的每像素预测结果,FCN引入了反卷积层(Deconvolution Layer) ,也称为转置卷积(Transposed Convolution) ,实现可学习的上采样操作。该过程可先以双线性插值 进行初始化,再通过训练进一步优化,使模型能够生成更加细致、准确的分割图。同时,为了弥补高层特征图中空间信息的缺失,FCN设计了跳跃连接(Skip Connections),将低层次(高分辨率)特征与高层次(高语义)特征进行融合,有效提升边缘细节的预测能力与目标定位精度,并在此基础上提出了FCN-16s与FCN-8s等更精细的改进版本。


二、FCN网络结构详解
FCN 网络的最大特点是:完全由卷积(convolution)、池化(pooling)、激活函数(ReLU)和上采样(反卷积)组成,不含全连接层 。它是从图像分类模型(如VGG、AlexNet)中"转化"而来,并重新设计用于像素级的语义分割任务。下面我们以论文中的 FCN-VGG16 为例,逐步分析其结构演变过程:
2.1 从分类网络到全卷积网络(Fully Convolutional)
FCN网络结构主要分为两个部分:全卷积部分 和反卷积(上采样)部分 。其中,全卷积部分由传统的图像分类网络(如 VGG16、ResNet 等)构成,用于逐层提取图像的语义特征;反卷积部分则负责将这些压缩后的语义特征图上采样还原为与输入图像相同大小的语义分割图。最终输出的每个像素点,代表它所属类别的概率分布。FCN的最大特点是打破了传统CNN只能接受固定尺寸输入的限制。通过移除全连接层(Fully Connected)并替换为等效的卷积操作,FCN能够接受任意尺寸的输入图像,并保持卷积神经网络的空间信息结构。
以经典的 VGG16 网络为例,其原始设计目标是用于图像分类任务,即判断整张图像属于哪一个类别。它的结构可以分为两个部分:
- 特征提取部分 :由 13个卷积层(conv) 和 5个最大池化层(max pooling) 组成,用于逐层提取图像的空间和语义特征。
- 分类决策部分 :由 3个全连接层(fc6、fc7、fc8) 组成,将前面提取到的高层特征压缩为一个固定维度的向量,最终输出图像所属的类别。
如下图和表为原始VGG16结构简要分布:
阶段 | 网络层(顺序) | 输出尺寸(输入224×224为例) |
---|---|---|
Block 1 | Conv1_1 → Conv1_2 → MaxPool | 112×112 |
Block 2 | Conv2_1 → Conv2_2 → MaxPool | 56×56 |
Block 3 | Conv3_1 → Conv3_2 → Conv3_3 → MaxPool | 28×28 |
Block 4 | Conv4_1 → Conv4_2 → Conv4_3 → MaxPool | 14×14 |
Block 5 | Conv5_1 → Conv5_2 → Conv5_3 → MaxPool | 7×7 |
Classifier | fc6(4096) → fc7(4096) → fc8(1000) | 1×1 |
注:fc8 的输出通常对应于 ImageNet 的1000个类别。

正是因为全连接层(Fully Connected Layer)本质上是将特征图展平(flatten)后进行矩阵乘法运算 ,它要求输入的特征图具有固定的空间尺寸,才能匹配预定义的权重维度。例如在 VGG16 中,输入图像必须是 224×224,经过一系列卷积和池化操作后得到的特征图大小为 7×7×512,会被展平为一个 25088 维的向量,再送入 fc6
处理,其对应的权重矩阵维度为 4096×25088,是事先写死的。因此,一旦输入图像尺寸发生变化,展平后的特征向量维度也会改变,导致无法与权重匹配,网络将因维度不一致而报错。而更关键的是,全连接层会完全打乱输入特征图的空间结构信息 ,也就是说,在进入 fc6
后,网络已无法感知哪些特征来自图像的哪个位置。这种结构虽然适合图像级别的分类任务,但对于需要保留像素空间位置信息的语义分割任务而言是致命缺陷,因为我们需要对每一个像素做出精确的类别判断。
因此为了实现网络对任意尺寸图像的处理能力,并保留空间结构以便输出每个像素的分类结果 ,FCN 对传统分类网络进行了结构性改造 ------ 即所谓的 "卷积化(Convolutionalize)" ,将原有的三个全连接层 fc6
、fc7
和 fc8
替换为尺寸等效的卷积层:
原始结构 | 卷积化后的替代层 | 说明 |
---|---|---|
fc6 |
conv6 (kernel=7×7) |
等效于对7×7感受野做全连接,输出4096通道 |
fc7 |
conv7 (kernel=1×1) |
提取语义特征 |
fc8 |
conv8 (kernel=1×1) |
输出每个空间位置上的类别分布(如21类) |

这样,网络的输出就变成了一张尺寸更小但仍保留空间结构的 score map(类别预测图),而非一个单一的分类向量。例如当我们输入图像尺寸为 512×512时经VGG16卷积+池化后输出为 16×16(下采样32倍),这时候每个位置输出一个长度为21的向量,表示该感受野区域对应的像素属于各类别的概率。这一结构上的转化,使得网络不仅可以处理任意尺寸图像,还能对图像中的每个位置进行分类预测,成为语义分割任务的基础。
2.2 特征图下采样与空间分辨率问题
在卷积神经网络(CNN)中,每一次卷积和池化操作都会对输入特征图进行空间下采样 ,即分辨率逐步减小。这种设计初衷是为了提取更加抽象的高级语义特征,同时减少计算量和内存占用。然而,对于语义分割这种需要像素级预测的任务来说,下采样过多会带来严重的空间信息丢失,尤其是在物体边缘区域,导致预测结果模糊不清。我们还是以FCN用到的VGG16来举例,以VGG16为例,如果输入图像的尺寸为512×512,经过VGG16网络的卷积和池化操作后,特征图的尺寸会逐步减小。在VGG16中,由于使用了5个池化层,每个池化层的步长为2,因此每经过一个池化层,特征图的尺寸就会缩小一半。最终,在经过最后一个池化层后,输入尺寸为512×512的图像,经过卷积和池化后的特征图尺寸将变为16×16。也就是说,特征图的空间尺寸会被下采样32倍(512/16=32)。
层级 | 特征图尺寸(H×W) | 下采样倍数 |
---|---|---|
输入 | 224×224 | 1× |
Conv1 → Pool1 | 112×112 | 2× |
Conv2 → Pool2 | 56×56 | 4× |
Conv3 → Pool3 | 28×28 | 8× |
Conv4 → Pool4 | 14×14 | 16× |
Conv5 → Pool5 | 7×7 | 32× |
也就是说,输入一张 512×512 的图像,经过 VGG16 后输出的特征图仅为 16×16,最终我们得到的语义特征图(即 conv5
输出)只有输入图尺寸的 1/32 大小 。意味着每个位置预测的结果实际上对应输入图上的一个 32×32 的区域(感受野即网络在该位置所看到的输入图像区域) ,这对于物体边缘或细小结构来说是非常粗糙的。由于语义分割的目标是:为图像中的每一个像素 分配一个语义标签。下采样过度的网络会导致输出的特征图空间尺寸过小,使得每个像素的预测实际上代表了输入图像上的一个较大区域。如果网络最后输出的特征图过小,我们只能得到非常稀疏的分类结果,哪怕后续再通过插值或上采样恢复图像尺寸,也会因为高频细节已经丢失而无法精确还原边界。

为了弥补下采样带来的空间精度损失,FCN提出了反卷积(Deconvolution) 或转置卷积(Transposed Convolution) 的机制。在网络的尾部,通过反卷积操作对特征图进行逐步的上采样,将低分辨率的特征图恢复到与原图相同的尺寸。反卷积操作能够在一定程度上将特征图恢复到原图的尺寸,但仅靠反卷积并不能完全恢复丢失的高频细节信息,尤其是在物体的边缘区域。为了解决这一问题,FCN进一步引入了**跳跃连接(Skip Connections)**的机制。跳跃连接通过将浅层特征(具有较高空间分辨率)与深层特征(包含较强语义信息)进行融合,有效地弥补了信息的丢失。通过这种方式,FCN能够在保持高层语义信息的同时,利用低层的细节特征来增强分割结果的精度,尤其是对于图像中的边缘和细小区域。这也就形成了后续我们要讲到的FCN-32s和FCN-16s以及FCN-8s

2.3 上采样(Upsampling):使用反卷积恢复原图尺寸
在FCN网络中,上采样(Upsampling)是一个关键步骤,它负责将经过多次下采样的低分辨率特征图恢复到与输入图像相同的尺寸,从而实现像素级的语义预测。FCN采用**反卷积(Deconvolution)**技术来实现这一过程,常见的上采样方法主要有三种:
方法 | 原理 | 优点 | 缺点 |
---|---|---|---|
最近邻插值 | 直接复制相邻像素值 | 计算简单,无参数 | 产生块状效应,质量差 |
双线性插值 | 通过线性加权平均进行平滑插值 | 效果较平滑 | 无法学习优化,细节恢复有限 |
反卷积/转置卷积 | 通过可学习的卷积核进行上采样 | 可训练优化,恢复效果最佳 | 计算量较大 |
FCN选择使用反卷积(Deconvolution) ,也称为转置卷积(Transposed Convolution) ,这种可学习的上采样方式相比传统插值方法具有显著优势。不同于数学上严格的逆卷积运算,反卷积实际上是通过在输入特征图元素间插入零值(通常插入stride-1
个零)并进行标准卷积操作来实现上采样。具体实现包含三个步骤:首先在空间维度进行零填充,然后在边缘补零(补零数量为kernel_size-padding-1
),最后使用转置后的卷积核执行常规卷积计算。这种设计不仅保留了卷积的参数共享特性,还能通过端到端训练自动学习最优的上采样方式,从而更有效地恢复特征图的空间细节。例如,当stride=2时,一个2×2的输入特征图经过零值插入后会扩展为3×3的矩阵,再通过卷积运算输出4×4的特征图,实现2倍上采样。这种可微分的上采样机制使FCN能够逐步重建高分辨率特征图,同时保持计算效率。反卷积的具体教学可以参考:反卷积(Transposed Convolution)详细推导 - 知乎

2.4 跳跃连接(Skip Connections):细节与语义的结合
FCN的创新性不仅体现在全卷积结构和上采样机制上,其**跳跃连接(Skip Connections)**的设计更是将语义分割的精度提升到了新的高度。这种架构灵感来源于人类视觉系统的多尺度信息整合能力------我们识别物体时既需要全局的语义理解,也需要局部的细节特征。在深度神经网络中,浅层特征往往包含丰富的空间细节(如边缘、纹理等),但由于感受野有限,语义理解能力较弱;而深层特征具有强大的语义表征能力,却因多次下采样丢失了空间细节。FCN通过跳跃连接创造性地解决了这一矛盾,实现了多尺度特征的有机融合。
具体实现上,FCN采用了一种金字塔式的特征融合策略 。以FCN-8s为例,网络首先将最深层的conv7特征进行2倍上采样,然后与来自pool4的同分辨率特征相加融合;接着对融合后的特征再次进行2倍上采样,与pool3的特征进行二次融合;最后通过8倍上采样得到最终预测结果。这种渐进式的融合方式犹如搭建金字塔,每一层都注入相应尺度的特征信息,使得网络在保持高层语义准确性的同时,能够精确恢复物体的边界细节。而在特征融合前需要对浅层特征进行1×1卷积处理,这既是为了调整通道维度,更是为了让不同层次的特征在语义空间中对齐,避免简单的特征堆叠导致优化困难。

从数学角度看,跳跃连接实际上构建了一个残差学习框架。假设最终预测结果为F(x),深层特征提供的基础预测为G(x),浅层特征提供的细节修正为H(x),则有F(x)=G(x)+H(x)。这种结构使网络更易于学习细节修正量,而不是直接学习复杂的映射关系,大大提升了训练效率和模型性能。实验数据显示,引入跳跃连接的FCN-8s在PASCAL VOC数据集上的mIoU达到62.7%,比没有跳跃连接的FCN-32s提高了近8个百分点,特别是在细小物体和复杂边界的分割上表现尤为突出。

三、实战搭建FCN模型
纸上得来终觉浅,绝知此事要躬行。理解了FCN的原理后,接着我将手把手带着大家用PyTorch从零开始搭建一个完整的FCN-8s模型,同时本项目已传至Github:xiongqi123123/Pytorch_FCN
3.1 环境准备与数据加载
我的开发环境如下:
WSL2-Ubuntu22.04
Python:3.10
PyTorch:2.0.1 Torchvision==0.15.2
GPU:NVIDIA GeForce RTX 3060 Cuda:12.5
我们首先配置我们的开发环境:
bash
# step1:创建一个Conda环境并激活
conda create -n FCN python=3.10 -y
conda activate FCN
# step2:下载安装依赖
pip install torch==2.0.1 torchvision==0.15.2
pip install numpy opencv-python matplotlib tqdm pillow
# step3:验证
python -c "import torch; print(torch.__version__, torch.cuda.is_available())"
验证之后如果终端显示True,即代表我们的Pytorch安装正确,如若遇到错误请自行搜索解决方法

数据集方面我们选择使用PASCAL VOC 2012,这是语义分割的经典基准数据集,大概有2GB大小,其下载方式如下,同时其解压后的目录格式大致如下:
bash
# step1:下载PASCAL VOC 2012数据集
wget http://host.robots.ox.ac.uk/pascal/VOC/voc2012/VOCtrainval_11-May-2012.tar
# step2:解压
tar -vxf VOCtrainval_11-May-2012.tar
# step3:验证目录结构
tree data/VOCdevkit/VOC2012 -L 2
VOCdevkit/
└── VOC2012/
├── Annotations/ # 目标检测标注(XML)
├── ImageSets/ # 数据集划分文件
│ └── Segmentation/ # 语义分割专用划分
├── JPEGImages/ # 原始图片
├── SegmentationClass/ # 类别标注图(PNG)
└── SegmentationObject/# 实例标注图(PNG)
3.1.1 DataLoad导入模块及宏变量
接下来我们来完成数据加载的部分 dataload.py
,这是训练和验证过程中不可或缺的一步。我们首先需要导入一些必要的模块,并预定义VOC数据集的类别名称 、类别数量 、图像重采样方式 以及用于可视化的**颜色映射表(colormap)**等宏变量
python
import os # 导入os模块,用于文件路径操作
import numpy as np # 导入numpy模块,用于数值计算
import torch # 导入PyTorch主模块
from torch.utils.data import Dataset, DataLoader # 导入数据集和数据加载器
from PIL import Image # 导入PIL图像处理库
import torchvision.transforms as transforms # 导入图像变换模块
from torchvision.transforms import functional as F # 导入函数式变换模块
# VOC数据集的类别名称(21个类别,包括背景)
VOC_CLASSES = [
'background', 'aeroplane', 'bicycle', 'bird', 'boat', 'bottle',
'bus', 'car', 'cat', 'chair', 'cow', 'diningtable', 'dog',
'horse', 'motorbike', 'person', 'potted plant', 'sheep', 'sofa',
'train', 'tv/monitor'
]
# 宏定义获取类别数量
NUM_CLASSES = len(VOC_CLASSES)
# 定义PIL的重采样常量
PIL_NEAREST = 0 # 最近邻重采样方式,保持锐利边缘,适用于掩码
PIL_BILINEAR = 1 # 双线性重采样方式,平滑图像,适用于原始图像
# 定义VOC数据集的颜色映射 (用于可视化分割结果),每个类别对应一个RGB颜色
VOC_COLORMAP = [
[0, 0, 0], [128, 0, 0], [0, 128, 0], [128, 128, 0], [0, 0, 128],
[128, 0, 128], [0, 128, 128], [128, 128, 128], [64, 0, 0], [192, 0, 0],
[64, 128, 0], [192, 128, 0], [64, 0, 128], [192, 0, 128], [64, 128, 128],
[192, 128, 128], [0, 64, 0], [128, 64, 0], [0, 192, 0], [128, 192, 0],
[0, 64, 128]
]
3.1.2 创建VOCSegmentation(Dataset)类
接着我们创建一个类 VOCSegmentation(Dataset)
,用于封装和加载 VOC2012 语义分割数据集。该类继承自 PyTorch 的 Dataset
,实现了标准的 __getitem__
和 __len__
方法,可直接配合 DataLoader
批量加载数据。它能够根据数据划分文件(如 train.txt
、val.txt
)读取图像与对应的掩码路径,并对它们进行统一尺寸的预处理------图像采用双线性插值以保持平滑性,掩码则使用最近邻插值以避免引入伪标签。同时,彩色掩码图会根据预定义的 VOC_COLORMAP
进行颜色到类别索引的映射,最终转换为模型训练所需的二维整数张量,实现从 RGB 掩码到语义标签的准确转换。
python
class VOCSegmentation(Dataset):
"""
VOC2012语义分割数据集的PyTorch Dataset实现
负责数据的加载、预处理和转换
"""
def __init__(self, root, split='train', transform=None, target_transform=None, img_size=320):
"""
初始化数据集
参数:
root (string): VOC数据集的根目录路径
split (string, optional): 使用的数据集划分,可选 'train', 'val' 或 'trainval'
transform (callable, optional): 输入图像的变换函数
target_transform (callable, optional): 目标掩码的变换函数
img_size (int, optional): 调整图像和掩码的大小
"""
def __getitem__(self, index):
"""
获取数据集中的一个样本
参数:
index (int): 样本索引
返回:
tuple: (图像, 掩码) 对,分别为图像张量和掩码张量
"""
def __len__(self):
"""返回数据集中的样本数量"""
我们来具体实现这三个类函数,首先是 __init__
。在该构造函数中,我们传入数据集的根目录 root
、划分类型 split
(如 'train'
、'val'
或 'trainval'
)、图像变换 transform
、标签变换 target_transform
以及目标尺寸 img_size
。随后根据划分文件(如 train.txt
)读取图像和掩码的文件名,拼接得到完整路径,并对路径有效性进行检查。最后,我们将图像和掩码路径分别保存在 self.images
和 self.masks
中,便于后续索引使用。
python
def __init__(self, root, split='train', transform=None, target_transform=None, img_size=320):
super(VOCSegmentation, self).__init__()
self.root = root
self.split = split
self.transform = transform
self.target_transform = target_transform
self.img_size = img_size
# 确定图像和标签文件的路径
voc_root = self.root
image_dir = os.path.join(voc_root, 'JPEGImages') # 原始图像目录
mask_dir = os.path.join(voc_root, 'SegmentationClass') # 语义分割标注目录
# 获取图像文件名列表(从划分文件中读取)
splits_dir = os.path.join(voc_root, 'ImageSets', 'Segmentation')
split_file = os.path.join(splits_dir, self.split + '.txt')
# 确保分割文件存在
if not os.path.exists(split_file):
raise FileNotFoundError(f"找不到拆分文件: {split_file}")
# 读取文件名列表
with open(split_file, 'r') as f:
file_names = [x.strip() for x in f.readlines()]
# 构建图像和掩码的完整路径
self.images = [os.path.join(image_dir, x + '.jpg') for x in file_names]
self.masks = [os.path.join(mask_dir, x + '.png') for x in file_names]
# 检查文件是否存在,打印警告但不中断程序
for img_path in self.images:
if not os.path.exists(img_path):
print(f"警告: 图像文件不存在: {img_path}")
for mask_path in self.masks:
if not os.path.exists(mask_path):
print(f"警告: 掩码文件不存在: {mask_path}")
# 确保图像和掩码数量匹配
assert len(self.images) == len(self.masks), "图像和掩码数量不匹配"
print(f"加载了 {len(self.images)} 对图像和掩码用于{split}集")
接下来是 __getitem__
方法,用于根据索引加载一个样本。图像和掩码被读取并转换为 RGB 格式,再统一调整为设定大小,其中图像使用双线性插值以保持平滑性,掩码使用最近邻插值以保留类别标签。图像经过指定的变换函数处理后,掩码则根据是否提供 target_transform
进行处理;若未指定,我们将掩码由 RGB 图转为类别索引图,通过遍历预定义的 VOC_COLORMAP
映射每个像素所属的语义类别,最终转为 long
类型的 PyTorch 张量,便于模型训练使用。最后的__len__
方法比较简单,直接返回数据集中图像的总数就好了,也就是 self.images
的长度,用于告知 PyTorch DataLoader
数据集的大小。
python
def __getitem__(self, index):
# 加载图像和掩码
img_path = self.images[index]
mask_path = self.masks[index]
img = Image.open(img_path).convert('RGB')# 打开图像并转换为RGB格式
mask = Image.open(mask_path).convert('RGB')# 打开掩码并转换为RGB格式(确保与colormap匹配)
# 统一调整图像和掩码大小,确保尺寸一致
img = img.resize((self.img_size, self.img_size), PIL_BILINEAR)# 对于图像使用双线性插值以保持平滑
mask = mask.resize((self.img_size, self.img_size), PIL_NEAREST)# 对于掩码使用最近邻插值以避免引入新的类别值
# 应用图像变换
if self.transform is not None:
img = self.transform(img)
# 处理掩码变换
if self.target_transform is not None:
mask = self.target_transform(mask)
else:
# 将掩码转换为类别索引
mask = np.array(mask)
# 检查掩码的维度,确保是RGB(3通道)
if len(mask.shape) != 3 or mask.shape[2] != 3:
raise ValueError(f"掩码维度错误: {mask.shape}, 期望为 (H,W,3)")
mask_copy = np.zeros((mask.shape[0], mask.shape[1]), dtype=np.uint8)# 创建一个新的类别索引掩码
# 将RGB颜色映射到类别索引
# 遍历每种颜色,将对应像素设置为类别索引
for k, color in enumerate(VOC_COLORMAP):
# 将每个颜色通道转换为布尔掩码
r_match = mask[:, :, 0] == color[0]
g_match = mask[:, :, 1] == color[1]
b_match = mask[:, :, 2] == color[2]
# 只有三个通道都匹配的像素才被分配为此类别
color_match = r_match & g_match & b_match
mask_copy[color_match] = k
mask = torch.from_numpy(mask_copy).long() # 转换为PyTorch张量(长整型,用于交叉熵损失)
return img, mask
def __len__(self):
return len(self.images)
因此完整的类代码如下:
python
class VOCSegmentation(Dataset):
def __init__(self, root, split='train', transform=None, target_transform=None, img_size=320):
super(VOCSegmentation, self).__init__()
self.root = root
self.split = split
self.transform = transform
self.target_transform = target_transform
self.img_size = img_size
# 确定图像和标签文件的路径
voc_root = self.root
image_dir = os.path.join(voc_root, 'JPEGImages') # 原始图像目录
mask_dir = os.path.join(voc_root, 'SegmentationClass') # 语义分割标注目录
# 获取图像文件名列表(从划分文件中读取)
splits_dir = os.path.join(voc_root, 'ImageSets', 'Segmentation')
split_file = os.path.join(splits_dir, self.split + '.txt')
# 确保分割文件存在
if not os.path.exists(split_file):
raise FileNotFoundError(f"找不到拆分文件: {split_file}")
# 读取文件名列表
with open(split_file, 'r') as f:
file_names = [x.strip() for x in f.readlines()]
# 构建图像和掩码的完整路径
self.images = [os.path.join(image_dir, x + '.jpg') for x in file_names]
self.masks = [os.path.join(mask_dir, x + '.png') for x in file_names]
# 检查文件是否存在,打印警告但不中断程序
for img_path in self.images:
if not os.path.exists(img_path):
print(f"警告: 图像文件不存在: {img_path}")
for mask_path in self.masks:
if not os.path.exists(mask_path):
print(f"警告: 掩码文件不存在: {mask_path}")
# 确保图像和掩码数量匹配
assert len(self.images) == len(self.masks), "图像和掩码数量不匹配"
print(f"加载了 {len(self.images)} 对图像和掩码用于{split}集")
def __getitem__(self, index):
# 加载图像和掩码
img_path = self.images[index]
mask_path = self.masks[index]
img = Image.open(img_path).convert('RGB')# 打开图像并转换为RGB格式
mask = Image.open(mask_path).convert('RGB')# 打开掩码并转换为RGB格式(确保与colormap匹配)
# 统一调整图像和掩码大小,确保尺寸一致
img = img.resize((self.img_size, self.img_size), PIL_BILINEAR)# 对于图像使用双线性插值以保持平滑
mask = mask.resize((self.img_size, self.img_size), PIL_NEAREST)# 对于掩码使用最近邻插值以避免引入新的类别值
# 应用图像变换
if self.transform is not None:
img = self.transform(img)
# 处理掩码变换
if self.target_transform is not None:
mask = self.target_transform(mask)
else:
# 将掩码转换为类别索引
mask = np.array(mask)
# 检查掩码的维度,确保是RGB(3通道)
if len(mask.shape) != 3 or mask.shape[2] != 3:
raise ValueError(f"掩码维度错误: {mask.shape}, 期望为 (H,W,3)")
mask_copy = np.zeros((mask.shape[0], mask.shape[1]), dtype=np.uint8)# 创建一个新的类别索引掩码
# 将RGB颜色映射到类别索引
# 遍历每种颜色,将对应像素设置为类别索引
for k, color in enumerate(VOC_COLORMAP):
# 将每个颜色通道转换为布尔掩码
r_match = mask[:, :, 0] == color[0]
g_match = mask[:, :, 1] == color[1]
b_match = mask[:, :, 2] == color[2]
# 只有三个通道都匹配的像素才被分配为此类别
color_match = r_match & g_match & b_match
mask_copy[color_match] = k
mask = torch.from_numpy(mask_copy).long() # 转换为PyTorch张量(长整型,用于交叉熵损失)
return img, mask
def __len__(self):
return len(self.images)
3.1.3 获取图像变换函数
在构建语义分割数据加载流程时,图像的预处理与增强变换同样不可或缺。我们定义了一个 get_transforms(train=True)
的辅助函数,根据当前阶段是否为训练集来决定变换策略。它返回一个二元组 (transform, target_transform)
,分别作用于输入图像和对应掩码。在训练阶段,我对输入图像加入了一系列增强策略以提升模型的泛化能力。例如:RandomHorizontalFlip()
:以概率0.5随机进行水平翻转,模拟真实场景中物体左右分布的多样性;ColorJitter()
:轻微扰动图像的亮度、对比度和饱和度,使模型能适应不同光照条件;ToTensor()
:将 PIL
图像转为 PyTorch 张量,并将像素值归一化到 [0, 1]
;Normalize()
:使用 ImageNet 数据集的均值和标准差对图像标准化,有利于预训练模型迁移。而在验证阶段,我们仅保留基本的 ToTensor()
和 Normalize()
,避免引入额外噪声,确保评估的客观性和稳定性。其中掩码图像无需归一化或张量化处理 ,因为我们在 __getitem__
中已将其转换为类别索引图。因此,target_transform
在此处设为 None
即可。
python
def get_transforms(train=True):
"""
获取图像变换函数
参数:
train (bool): 是否为训练集,决定是否应用数据增强
返回:
tuple: (图像变换, 目标掩码变换)
"""
if train:
# 训练集使用数据增强
transform = transforms.Compose([
transforms.RandomHorizontalFlip(),# 随机水平翻转增加数据多样性
transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),# 随机调整亮度、对比度和饱和度
transforms.ToTensor(),# 转换为张量(值范围变为[0,1])
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) # 归一化(使用ImageNet的均值和标准差)
])
else:
# 验证集只需要基本变换
transform = transforms.Compose([
transforms.ToTensor(), # 转换为张量
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])# 归一化
])
target_transform = None# 掩码不需要标准化或者转换为Tensor (已在__getitem__中处理)
return transform, target_transform
3.1.4 创建训练和验证数据加载器
数据集和变换函数准备好之后,接下来我们通过 get_data_loaders
函数统一创建训练和验证的数据加载器 train_loader
与 val_loader
。这个函数首先将我们封装好的 VOCSegmentation
数据集类实例化,接着设置 DataLoader
中的一些参数,如 batch_size
、线程数量 num_workers
、是否打乱 shuffle
等,方便后续模型训练与验证阶段高效批量读取数据。
对于训练集,我们启用了 shuffle=True
以打乱数据顺序,增强模型的泛化能力,同时开启 drop_last=True
来舍弃最后一个不足 batch 的数据,避免 batchnorm 等层出现异常。而验证集则保持顺序读取(shuffle=False
),确保评估过程的稳定性。同时我们开启了pin_memory=True
参数,这个参数能够将数据预加载到锁页内存中,加快从 CPU 到 GPU 的数据拷贝效率。
该函数最终返回构建好的训练与验证加载器,可直接用于训练循环中的迭代操作,完整代码如下所示:
python
def get_data_loaders(voc_root, batch_size=4, num_workers=4, img_size=320):
"""
创建训练和验证数据加载器
参数:
voc_root (string): VOC数据集根目录
batch_size (int): 批次大小
num_workers (int): 数据加载的线程数
img_size (int): 图像的大小
返回:
tuple: (train_loader, val_loader) 训练和验证数据加载器
"""
# 获取图像和掩码变换
train_transform, train_target_transform = get_transforms(train=True)
val_transform, val_target_transform = get_transforms(train=False)
# 创建训练数据集
train_dataset = VOCSegmentation(
root=voc_root,
split='train', # 使用训练集划分
transform=train_transform,
target_transform=train_target_transform,
img_size=img_size
)
# 创建验证数据集
val_dataset = VOCSegmentation(
root=voc_root,
split='val', # 使用验证集划分
transform=val_transform,
target_transform=val_target_transform,
img_size=img_size
)
# 创建训练数据加载器
train_loader = DataLoader(
train_dataset,
batch_size=batch_size,
shuffle=True, # 随机打乱数据
num_workers=num_workers, # 多线程加载
pin_memory=True, # 数据预加载到固定内存,加速GPU传输
drop_last=True # 丢弃最后不足一个批次的数据
)
# 创建验证数据加载器
val_loader = DataLoader(
val_dataset,
batch_size=batch_size,
shuffle=False, # 不打乱数据
num_workers=num_workers,
pin_memory=True
)
return train_loader, val_loader
3.1.5 可视化分割结果
在训练语义分割模型的过程中,如果我们只是单纯地关注每轮训练的数值指标(如 IoU、准确率等),难免会显得有些枯燥,且难以直观感受模型到底学得怎么样 。尤其是在模型逐步收敛时,仅靠指标的波动并不能很好地揭示模型的细节表现。因此,我在此基础上引入了一个可视化辅助函数 decode_segmap
,用于将模型预测得到的分割结果从类别索引图转换为彩色图像。这样一来,我们就可以将每个像素所属的类别清晰地呈现在图像上,借助这个工具,我们可以在训练过程中插入实时可视化,随时查看模型对于不同样本的分割表现,为调参和模型改进提供更加直观的反馈。完整实现的代码如下:
python
def decode_segmap(segmap):
"""
将类别索引的分割图转换为RGB彩色图像(用于可视化)
参数:
segmap (np.array或torch.Tensor): 形状为(H,W)的分割图,值为类别索引
返回:
rgb_img (np.array): 形状为(H,W,3)的RGB彩色图像
"""
# 确保segmap是NumPy数组
if isinstance(segmap, torch.Tensor):
segmap = segmap.cpu().numpy()
# 检查segmap的形状,处理各种可能的输入格式
if len(segmap.shape) > 2:
if len(segmap.shape) == 3 and segmap.shape[0] <= 3:
segmap = segmap[0]
rgb_img = np.zeros((segmap.shape[0], segmap.shape[1], 3), dtype=np.uint8) # 创建RGB图像
# 根据类别索引填充对应的颜色
for cls_idx, color in enumerate(VOC_COLORMAP):
mask = segmap == cls_idx# 找到属于当前类别的像素
if mask.any(): # 只处理存在的类别
rgb_img[mask] = color # 将这些像素设置为对应的颜色
return rgb_img
3.2 FCN主干网络搭建
在完成了数据加载函数之后我们便可以开始完成FCN主干网络的搭建了,我们将从最基础的FCN32s主干网络开始,逐步完成FCN16s和FCN8s,首先我们先创建fcn_model.py
文件,同时导入一些必要的模块
python
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.models as models
from dataload import NUM_CLASSES
3.2.1 完成双线性插值辅助函数
在构建FCN网络时,我们需要用到反卷积(ConvTranspose2d
)层来进行上采样,而这种上采样常常使用双线性插值初始化,以确保上采样时图像的平滑性。我们将在网络的_initialize_weights
方法中使用我们自定义的双线性插值方法来初始化这些反卷积层的权重,。
python
def _initialize_weights(self):
# 初始化反卷积层的权重为双线性上采样
for m in self.modules():
if isinstance(m, nn.ConvTranspose2d):
# 双线性上采样的初始化
m.weight.data.zero_()
m.weight.data = self._make_bilinear_weights(m.kernel_size[0], m.out_channels)
上面的代码确保了反卷积层的权重被正确初始化为适合图像上采样的双线性插值权重。接下来,我们实现_make_bilinear_weights
方法用于生成双线性插值的权重矩阵。
python
def _make_bilinear_weights(self, size, num_channels):
"""生成双线性插值的权重"""
factor = (size + 1) // 2
if size % 2 == 1:
center = factor - 1
else:
center = factor - 0.5
og = torch.FloatTensor(size, size)
for i in range(size):
for j in range(size):
og[i, j] = (1 - abs((i - center) / factor)) * (1 - abs((j - center) / factor))
filter = torch.zeros(num_channels, num_channels, size, size)
for i in range(num_channels):
filter[i, i] = og
return filter
通过这个方法,我们可以生成一个num_channels
个通道的双线性插值权重矩阵,并将其赋值给反卷积层的权重。这样,网络在训练时,反卷积层将能够根据这些初始化的权重执行平滑的上采样操作。
3.2.2 搭建FCN32s主干网络
FCN32s是FCN系列中最基础的版本,其核心思想是直接将VGG16的全连接层转换为卷积层,并通过32倍上采样恢复原图分辨率。下面我将详细讲解如何用PyTorch实现这一结构,首先我们创建class FCN32s(nn.Module):
网络
python
class FCN32s(nn.Module):
def __init__(self, num_classes=NUM_CLASSES, pretrained=True):
super(FCN32s, self).__init__()
FCN的特征提取主干网络使用的是VGG16,torchvision中有现成的网络架构可以直接导入,但是由于我们要实现分割,因此我们导入了VGG16后还需要将VGG16网络特征层做一些修改,使其能够适应全卷积网络结构。主要操作便是使用Sequential
操作来提取VGG16中的卷积层。我们使用nn.Sequential
将VGG16的前五个卷积块封装保留下来,接着我们将VGG16网络中的全连接层替换为1x1卷积层,最后我们通过一个上采样操作将低分辨率的特征图恢复到原始输入图像的尺寸,从而能够进行像素级别的分割预测。
python
vgg16 = models.vgg16(pretrained=pretrained)# 加载预训练的VGG16模型
features = list(vgg16.features.children())# 获取特征提取部分
# 根据FCN原始论文修改VGG16网络
# 前5段卷积块保持不变
self.features1 = nn.Sequential(*features[:5]) # conv1 + pool1
self.features2 = nn.Sequential(*features[5:10]) # conv2 + pool2
self.features3 = nn.Sequential(*features[10:17]) # conv3 + pool3
self.features4 = nn.Sequential(*features[17:24]) # conv4 + pool4
self.features5 = nn.Sequential(*features[24:31]) # conv5 + pool5
# 全连接层替换为1x1卷积
self.fc6 = nn.Conv2d(512, 4096, kernel_size=7, padding=3)
self.relu6 = nn.ReLU(inplace=True)
self.drop6 = nn.Dropout2d()
self.fc7 = nn.Conv2d(4096, 4096, kernel_size=1)
self.relu7 = nn.ReLU(inplace=True)
self.drop7 = nn.Dropout2d()
# 分类层
self.score = nn.Conv2d(4096, num_classes, kernel_size=1)
# 上采样层: 32倍上采样回原始图像大小
self.upsample = nn.ConvTranspose2d(num_classes, num_classes, kernel_size=64, stride=32, padding=16, bias=False)
# 初始化参数
self._initialize_weights()
接着我们来实现这个网络的前向推理部分,在这部分里,我们首先记录输入图像的尺寸,input_size = x.size()[2:]
,这是为了在经过上采样后确保输出的尺寸与输入图像一致。接下来,输入图像会依次通过VGG16网络中的卷积层进行特征提取。我们将图像通过features1
到features5
(分别对应VGG16中的五个卷积块)进行处理之后,我们会得到了一个较低分辨率的特征图,为了将这些低分辨率特征图转化为像素级的预测,我们接下来将它们通过两层1x1卷积(即fc6
和fc7
)进行处理,并使用ReLU激活函数进行非线性转换,同时为了防止过拟合我们在每一层后都应用了Dropout
,接着经过分类层score
后,我们便得到了一个最终的输出特征图,其中每个像素点的通道对应于一个类别的分割结果,然后我们便可以通过转置卷积(upsample
)对输出进行32倍上采样,将特征图恢复到原始图像的尺寸。最后,我们裁剪输出的尺寸,确保它与输入图像的大小一致。
python
def forward(self, x):
input_size = x.size()[2:]# 记录输入尺寸用于上采样
# 编码器 (VGG16)
x = self.features1(x)
x = self.features2(x)
x = self.features3(x)
x = self.features4(x)
x = self.features5(x)
# 全连接层 (以卷积形式实现)
x = self.relu6(self.fc6(x))
x = self.drop6(x)
x = self.relu7(self.fc7(x))
x = self.drop7(x)
x = self.score(x)# 分类
x = self.upsample(x)# 上采样回原始尺寸
x = x[:, :, :input_size[0], :input_size[1]]# 裁剪到原始图像尺寸
return x
完整的FCN32s网络如下:
python
class FCN32s(nn.Module):
def __init__(self, num_classes=NUM_CLASSES, pretrained=True):
super(FCN32s, self).__init__()
vgg16 = models.vgg16(pretrained=pretrained)# 加载预训练的VGG16模型
features = list(vgg16.features.children())# 获取特征提取部分
# 根据FCN原始论文修改VGG16网络
# 前5段卷积块保持不变
self.features1 = nn.Sequential(*features[:5]) # conv1 + pool1
self.features2 = nn.Sequential(*features[5:10]) # conv2 + pool2
self.features3 = nn.Sequential(*features[10:17]) # conv3 + pool3
self.features4 = nn.Sequential(*features[17:24]) # conv4 + pool4
self.features5 = nn.Sequential(*features[24:31]) # conv5 + pool5
# 全连接层替换为1x1卷积
self.fc6 = nn.Conv2d(512, 4096, kernel_size=7, padding=3)
self.relu6 = nn.ReLU(inplace=True)
self.drop6 = nn.Dropout2d()
self.fc7 = nn.Conv2d(4096, 4096, kernel_size=1)
self.relu7 = nn.ReLU(inplace=True)
self.drop7 = nn.Dropout2d()
self.score = nn.Conv2d(4096, num_classes, kernel_size=1)# 分类层
# 上采样层: 32倍上采样回原始图像大小
self.upsample = nn.ConvTranspose2d(num_classes, num_classes, kernel_size=64, stride=32, padding=16, bias=False)
self._initialize_weights()# 初始化参数
def forward(self, x):
input_size = x.size()[2:]# 记录输入尺寸用于上采样
# 编码器 (VGG16)
x = self.features1(x)
x = self.features2(x)
x = self.features3(x)
x = self.features4(x)
x = self.features5(x)
# 全连接层 (以卷积形式实现)
x = self.relu6(self.fc6(x))
x = self.drop6(x)
x = self.relu7(self.fc7(x))
x = self.drop7(x)
x = self.score(x)# 分类
x = self.upsample(x)# 上采样回原始尺寸
x = x[:, :, :input_size[0], :input_size[1]]# 裁剪到原始图像尺寸
return x
def _initialize_weights(self):
# 初始化反卷积层的权重为双线性上采样
for m in self.modules():
if isinstance(m, nn.ConvTranspose2d):
# 双线性上采样的初始化
m.weight.data.zero_()
m.weight.data = self._make_bilinear_weights(m.kernel_size[0], m.out_channels)
def _make_bilinear_weights(self, size, num_channels):
"""生成双线性插值的权重"""
factor = (size + 1) // 2
if size % 2 == 1:
center = factor - 1
else:
center = factor - 0.5
og = torch.FloatTensor(size, size)
for i in range(size):
for j in range(size):
og[i, j] = (1 - abs((i - center) / factor)) * (1 - abs((j - center) / factor))
filter = torch.zeros(num_channels, num_channels, size, size)
for i in range(num_channels):
filter[i, i] = og
return filter
3.2.3 搭建FCN16s与FCN8s主干网络
FCN16s与FCN8s相比于FCN32s主要的变动便是引入了更多层级的特征图进行融合,从而提升分割结果的细节还原能力 。FCN32s的上采样仅使用了VGG16最后一个卷积块(conv5)后的输出,而FCN16s在此基础上引入了pool4的特征图 ,而FCN8s则进一步引入了pool3的特征图 ,这种特征融合策略可以有效提升空间细节的恢复效果。从网络的实现上来看FCN16s在执行最后一次上采样之前,会先对conv5
的输出进行上采样(2倍),然后与pool4
对应的特征图进行逐像素相加,接着再执行进一步上采样至原始尺寸。而FCN8s则在FCN16s的基础上再上采样2倍后与pool3
的特征图进行融合,最后再上采样8倍回到原图大小。我在这里只展示相比于FCN32s有变化的地方:
python
### FCN16s
class FCN16s(nn.Module):
def __init__(self, num_classes=NUM_CLASSES, pretrained=True):
# 获取特征提取部分
# 分段处理VGG16特征
######以上和FCN32s保持一致#########
# 全连接层替换为1x1卷积
self.fc6 = nn.Conv2d(512, 4096, kernel_size=7, padding=3)
self.relu6 = nn.ReLU(inplace=True)
self.drop6 = nn.Dropout2d()
self.fc7 = nn.Conv2d(4096, 4096, kernel_size=1)
self.relu7 = nn.ReLU(inplace=True)
self.drop7 = nn.Dropout2d()
self.score_fr = nn.Conv2d(4096, num_classes, kernel_size=1)# 分类层
self.score_pool4 = nn.Conv2d(512, num_classes, kernel_size=1)# pool4的1x1卷积,用于特征融合
# 2倍上采样conv7特征
self.upsample2 = nn.ConvTranspose2d(num_classes, num_classes, kernel_size=4, stride=2, padding=1, bias=False)
# 16倍上采样回原始图像大小
self.upsample16 = nn.ConvTranspose2d(num_classes, num_classes, kernel_size=32, stride=16, padding=8, bias=False)
# 初始化参数
self._initialize_weights()
def forward(self, x):
input_size = x.size()[2:]# 记录输入尺寸用于上采样
# 编码器 (VGG16)
x = self.features1(x)
x = self.features2(x)
x = self.features3(x)
pool4 = self.features4(x)# 保存pool4的输出用于后续融合
x = self.features5(pool4)
# 全连接层 (以卷积形式实现)
x = self.relu6(self.fc6(x))
x = self.drop6(x)
x = self.relu7(self.fc7(x))
x = self.drop7(x)
x = self.score_fr(x)# 分类
# 2倍上采样
x = self.upsample2(x)
# 获取pool4的分数并裁剪
score_pool4 = self.score_pool4(pool4)
score_pool4 = score_pool4[:, :, :x.size()[2], :x.size()[3]]
x = x + score_pool4# 融合特征
x = self.upsample16(x)# 16倍上采样回原始尺寸
x = x[:, :, :input_size[0], :input_size[1]]# 裁剪到原始图像尺寸
return x
### FCN8s
class FCN8s(nn.Module):
def __init__(self, num_classes=NUM_CLASSES, pretrained=True):
# 获取特征提取部分
# 分段处理VGG16特征
######以上和FCN32s保持一致#########
# 全连接层替换为1x1卷积
self.fc6 = nn.Conv2d(512, 4096, kernel_size=7, padding=3)
self.relu6 = nn.ReLU(inplace=True)
self.drop6 = nn.Dropout2d()
self.fc7 = nn.Conv2d(4096, 4096, kernel_size=1)
self.relu7 = nn.ReLU(inplace=True)
self.drop7 = nn.Dropout2d()
self.score_fr = nn.Conv2d(4096, num_classes, kernel_size=1)# 分类层
# pool3和pool4的1x1卷积,用于特征融合
self.score_pool4 = nn.Conv2d(512, num_classes, kernel_size=1)
self.score_pool3 = nn.Conv2d(256, num_classes, kernel_size=1)
# 2倍上采样conv7特征
self.upsample2_1 = nn.ConvTranspose2d(num_classes, num_classes, kernel_size=4, stride=2, padding=1, bias=False)
# 2倍上采样融合后的特征
self.upsample2_2 = nn.ConvTranspose2d(num_classes, num_classes, kernel_size=4, stride=2, padding=1, bias=False)
# 8倍上采样回原始图像大小
self.upsample8 = nn.ConvTranspose2d(num_classes, num_classes, kernel_size=16, stride=8, padding=4, bias=False)
# 初始化参数
self._initialize_weights()
def forward(self, x):
input_size = x.size()[2:]# 记录输入尺寸用于上采样
# 编码器 (VGG16)
x = self.features1(x)
x = self.features2(x)
pool3 = self.features3(x)# 保存pool3的输出用于后续融合
pool4 = self.features4(pool3)# 保存pool4的输出用于后续融合
x = self.features5(pool4)
# 全连接层 (以卷积形式实现)
x = self.relu6(self.fc6(x))
x = self.drop6(x)
x = self.relu7(self.fc7(x))
x = self.drop7(x)
x = self.score_fr(x)# 分类
x = self.upsample2_1(x)# 2倍上采样
# 获取pool4的分数并裁剪
score_pool4 = self.score_pool4(pool4)
score_pool4 = score_pool4[:, :, :x.size()[2], :x.size()[3]]
x = x + score_pool4 # 第一次融合特征 (pool5上采样 + pool4)
x = self.upsample2_2(x)# 再次2倍上采样
# 获取pool3的分数并裁剪
score_pool3 = self.score_pool3(pool3)
score_pool3 = score_pool3[:, :, :x.size()[2], :x.size()[3]]
x = x + score_pool3# 第二次融合特征 (第一次融合的上采样 + pool3)
x = self.upsample8(x)# 8倍上采样回原始尺寸
x = x[:, :, :input_size[0], :input_size[1]]# 裁剪到原始图像尺寸
return x
3.3 完成训练脚本train.py
在完成了FCN模型的网络结构搭建后,我们需要编写一个完整的训练脚本来对模型进行训练和评估。这个脚本将包含数据加载、模型训练、验证评估以及结果可视化等功能。下面我将详细讲解训练脚本的各个组成部分。
3.3.1 导入所需模块和解析命令行参数
首先,我们需要导入必要的模块,并定义一个参数解析器,使得我们可以通过命令行灵活地调整训练参数,在这里我们定义了一系列参数,包括数据集路径、模型类型选择、训练超参数(批大小、轮数、学习率等)以及检查点相关参数,使得我们可以灵活控制训练过程。
python
import os
import time
import argparse
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from tqdm import tqdm
import matplotlib.pyplot as plt
import gc
from dataload import get_data_loaders, NUM_CLASSES, decode_segmap
from fcn_model import get_fcn_model
def parse_args():
parser = argparse.ArgumentParser(description='FCN 语义分割 PyTorch 实现')
parser.add_argument('--voc-root', type=str, default='',
help='VOC数据集根目录')
parser.add_argument('--model-type', type=str, default='fcn8s', choices=['fcn8s', 'fcn16s', 'fcn32s'],
help='FCN模型类型 (fcn8s, fcn16s, fcn32s)')
parser.add_argument('--batch-size', type=int, default=4,
help='训练的批次大小')
parser.add_argument('--epochs', type=int, default=50,
help='训练的轮数')
parser.add_argument('--lr', type=float, default=0.005,
help='学习率')
parser.add_argument('--momentum', type=float, default=0.9,
help='SGD动量')
parser.add_argument('--weight-decay', type=float, default=1e-4,
help='权重衰减')
parser.add_argument('--num-workers', type=int, default=4,
help='数据加载线程数')
parser.add_argument('--checkpoint-dir', type=str, default='checkpoints',
help='模型保存目录')
parser.add_argument('--resume', type=str, default=None,
help='恢复训练的检查点路径')
return parser.parse_args()
3.3.2 模型评估函数
接下来,我们定义一个评估函数,用于在验证集上评估模型性能。语义分割任务常用的评估指标包括像素准确率和平均交并比(mIoU),在这个评估函数中,我们通过遍历验证集的每个批次,计算模型的预测结果与真实标签之间的损失、像素准确率以及每个类别的IoU(交并比)
python
def evaluate(model, val_loader, criterion, device):
model.eval()
total_loss = 0.0
total_corrects = 0
total_pixels = 0
class_iou = np.zeros(NUM_CLASSES)
class_pixels = np.zeros(NUM_CLASSES)
with torch.no_grad():
for images, targets in tqdm(val_loader, desc='Evaluation'):
images = images.to(device)
targets = targets.to(device)
outputs = model(images)
loss = criterion(outputs, targets)
total_loss += loss.item() * images.size(0)
_, preds = torch.max(outputs, 1)
# 计算像素准确率
correct = (preds == targets).sum().item()
total_corrects += correct
total_pixels += targets.numel()
# 计算每个类别的IoU
for cls in range(NUM_CLASSES):
pred_inds = preds == cls
target_inds = targets == cls
intersection = (pred_inds & target_inds).sum().item()
union = (pred_inds | target_inds).sum().item()
if union > 0:
class_iou[cls] += intersection / union
class_pixels[cls] += 1
del images, targets, outputs, preds
torch.cuda.empty_cache()
# 计算平均指标
val_dataset_size = len(val_loader.dataset) if hasattr(val_loader.dataset, '__len__') else len(val_loader) * val_loader.batch_size
avg_loss = total_loss / val_dataset_size
pixel_acc = total_corrects / total_pixels
# 计算每个类别的平均IoU
for cls in range(NUM_CLASSES):
if class_pixels[cls] > 0:
class_iou[cls] /= class_pixels[cls]
# 计算mIoU (平均交并比)
miou = np.mean(class_iou)
gc.collect()
torch.cuda.empty_cache()
return avg_loss, pixel_acc, miou, class_iou
3.3.3 分割结果可视化函数
为了直观地观察分割结果,我们定义一个函数用于保存预测结果的可视化图像。这个函数将从验证集中取出一定数量的样本,通过模型进行预测,然后将原始图像、真实标签和预测结果并排显示并保存为图像文件,这样我们就可以直观地观察模型的分割效果
python
def save_predictions(model, val_loader, device, output_dir='outputs', num_samples=5):
if not os.path.exists(output_dir):
os.makedirs(output_dir)
model.eval()
with torch.no_grad():
for i, (images, targets) in enumerate(val_loader):
if i >= num_samples:
break
images = images.to(device)
targets = targets.to(device)
outputs = model(images)
_, preds = torch.max(outputs, 1)
# 转换为NumPy数组用于可视化
images_np = images.cpu().numpy()
targets_np = targets.cpu().numpy()
preds_np = preds.cpu().numpy()
# 对每个样本进行可视化
for b in range(images.size(0)):
if b >= 3: # 限制每个批次只保存前3个样本
break
img = images_np[b].transpose(1, 2, 0)
mean = np.array([0.485, 0.456, 0.406])
std = np.array([0.229, 0.224, 0.225])
img = img * std + mean
img = np.clip(img, 0, 1)
target_rgb = decode_segmap(targets_np[b])
pred_rgb = decode_segmap(preds_np[b])
plt.figure(figsize=(15, 5))
plt.subplot(1, 3, 1)
plt.title('Input Image')
plt.imshow(img)
plt.axis('off')
plt.subplot(1, 3, 2)
plt.title('Ground Truth')
plt.imshow(target_rgb)
plt.axis('off')
plt.subplot(1, 3, 3)
plt.title('Prediction')
plt.imshow(pred_rgb)
plt.axis('off')
plt.tight_layout()
plt.savefig(os.path.join(output_dir, f'sample_{i}_{b}.png'))
plt.close()
del images, targets, outputs, preds
torch.cuda.empty_cache()
gc.collect()
torch.cuda.empty_cache()
3.3.4 主训练函数
最后,我们编写主函数,实现完整的训练流程。主函数是整个训练脚本的核心,它将各个组件有机地整合在一起,形成完整的训练流程。首先,它通过parse_args()
解析命令行输入的各项参数,如数据集路径、模型类型、批量大小等,使训练过程更加灵活可控。之后,它会调用get_data_loaders()
函数加载并预处理VOC数据集,同时创建数据加载器以便批量获取训练和验证样本。接着,根据参入参数指定的模型类型(FCN8s/FCN16s/FCN32s)实例化相应的网络结构,并将其转移到可用的计算设备(GPU或CPU)上。在优化策略方面,主函数使用交叉熵损失函数(忽略255标签值)评估分割质量,采用带动量的SGD优化器更新网络参数,并通过学习率调度器在训练后期降低学习率以获得更精细的优化效果。如果参入参数提供了检查点路径,函数会从中恢复模型权重、优化器状态和训练进度,实现断点续训。核心的训练循环涵盖了完整的训练-评估-保存流程:每个epoch内先在训练集上进行前向传播、损失计算、反向传播和参数更新;然后在验证集上评估模型性能(损失值、像素准确率和mIoU);当取得更高mIoU时,保存最佳模型并生成可视化结果,同时定期保存最新模型以防训练中断。训练完成后,主函数会绘制整个训练过程的损失曲线、准确率曲线和mIoU曲线,直观展示模型的学习轨迹和性能变化,帮助大家更好地理解训练动态并优化训练策略。
python
def main():
args = parse_args()
if not os.path.exists(args.checkpoint_dir):
os.makedirs(args.checkpoint_dir)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f'使用设备: {device}')
# 加载数据
train_loader, val_loader = get_data_loaders(
args.voc_root,
batch_size=args.batch_size,
num_workers=args.num_workers
)
train_dataset_size = len(train_loader.dataset) if hasattr(train_loader.dataset, '__len__') else len(train_loader) * train_loader.batch_size
val_dataset_size = len(val_loader.dataset) if hasattr(val_loader.dataset, '__len__') else len(val_loader) * val_loader.batch_size
print(f'训练样本数: {train_dataset_size}, 验证样本数: {val_dataset_size}')
# 创建模型
model = get_fcn_model(model_type=args.model_type, num_classes=NUM_CLASSES, pretrained=True)
model = model.to(device)
# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss(ignore_index=255)
optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)
# 恢复训练
start_epoch = 0
best_miou = 0.0
if args.resume:
if os.path.isfile(args.resume):
print(f'加载检查点: {args.resume}')
checkpoint = torch.load(args.resume)
start_epoch = checkpoint['epoch']
best_miou = checkpoint['best_miou']
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
print(f'从 epoch {start_epoch} 恢复训练, 最佳 mIoU: {best_miou:.4f}')
else:
print(f'找不到检查点: {args.resume}')
# 训练循环
print(f'开始训练 {args.model_type} 模型, 共 {args.epochs} 轮...')
# 记录训练历史
history = {
'train_loss': [],
'val_loss': [],
'pixel_acc': [],
'miou': []
}
for epoch in range(start_epoch, args.epochs):
# 训练阶段
model.train()
train_loss = 0.0
batch_count = 0
t0 = time.time()
for images, targets in tqdm(train_loader, desc=f'Epoch {epoch+1}/{args.epochs}'):
images = images.to(device)
targets = targets.to(device)
optimizer.zero_grad()
outputs = model(images)
loss = criterion(outputs, targets)
loss.backward()
optimizer.step()
train_loss += loss.item() * images.size(0)
batch_count += 1
del images, targets, outputs, loss
if batch_count % 10 == 0:
torch.cuda.empty_cache()
train_loss = train_loss / train_dataset_size
history['train_loss'].append(train_loss)
# 调整学习率
scheduler.step()
gc.collect()
torch.cuda.empty_cache()
# 评估模型
val_loss, pixel_acc, miou, class_iou = evaluate(model, val_loader, criterion, device)
history['val_loss'].append(val_loss)
history['pixel_acc'].append(pixel_acc)
history['miou'].append(miou)
# 打印进度
epoch_time = time.time() - t0
print(f'Epoch {epoch+1}/{args.epochs} - '
f'Time: {epoch_time:.2f}s - '
f'Train Loss: {train_loss:.4f} - '
f'Val Loss: {val_loss:.4f} - '
f'Pixel Acc: {pixel_acc:.4f} - '
f'mIoU: {miou:.4f}')
# 保存最佳模型
if miou > best_miou:
best_miou = miou
torch.save({
'epoch': epoch + 1,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'scheduler_state_dict': scheduler.state_dict(),
'best_miou': best_miou,
}, os.path.join(args.checkpoint_dir, f'{args.model_type}_best.pth'))
print(f'保存最佳模型, mIoU: {best_miou:.4f}')
# 生成可视化结果
save_predictions(model, val_loader, device)
# 保存最新模型
torch.save({
'epoch': epoch + 1,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'scheduler_state_dict': scheduler.state_dict(),
'best_miou': best_miou,
}, os.path.join(args.checkpoint_dir, f'{args.model_type}_latest.pth'))
gc.collect()
torch.cuda.empty_cache()
plt.figure(figsize=(12, 10))
plt.subplot(2, 2, 1)
plt.plot(history['train_loss'], label='Train')
plt.plot(history['val_loss'], label='Validation')
plt.title('Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()
plt.subplot(2, 2, 2)
plt.plot(history['pixel_acc'])
plt.title('Pixel Accuracy')
plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.subplot(2, 2, 3)
plt.plot(history['miou'])
plt.title('Mean IoU')
plt.xlabel('Epoch')
plt.ylabel('mIoU')
plt.tight_layout()
plt.savefig(os.path.join(args.checkpoint_dir, f'{args.model_type}_history.png'))
plt.close()
print(f'训练完成! 最佳 mIoU: {best_miou:.4f}')
完成了训练代码之后我们便可以开始训练啦!我们输入以下命令即可开始训练:
bash
python train.py --epochs 50 --batch-size 4 --lr 0.005 --model-type fcn8s

同时我们可以看到训练过程中我们的项目目录生成了两个文件夹,checkpoints
用于保存模型的最佳权重以及最后一次训练的权重,outputs
用于在训练过程中实时查看到我们的可视化训练分割结果

训练了22epoch后的结果如下,可以看到还有待进一步训练,mIoU:目前还只有0.2855



3.4 完成推理预测脚本predict.py
训练好模型后,我们需要一个单独的脚本来对新图像进行语义分割预测。这个推理脚本不仅能够加载我们训练好的模型,还能对单张图像或整个文件夹的图像进行批量预测,同时提供多种可视化方式展示分割结果。下面我将详细讲解这个推理脚本的实现过程。
3.4.1 导入必要模块和解析命令行参数
首先,我们需要导入必要的模块,并设置命令行参数解析器,以便灵活地配置推理过程。在参数解析部分,我们可以通过--model-path
指定预训练模型的存储路径;同时通过--model-type
选择使用FCN8s、FCN16s或FCN32s中的任一模型架构。而--image-path
参数支持单个图像文件也可以制定一个文件夹进行批量处理。分割结果默认保存在名为"results"的文件夹中,也可以通过--output-dir
参数自定义存储位置,--overlay
参数则可以选择是否将掩码叠加在原图上面
python
import os
import argparse
import numpy as np
import torch
import torchvision.transforms as transforms
from PIL import Image
import matplotlib.pyplot as plt
from dataload import VOC_CLASSES, VOC_COLORMAP, NUM_CLASSES, decode_segmap
from fcn_model import get_fcn_model
def parse_args():
parser = argparse.ArgumentParser(description='FCN语义分割模型预测')
parser.add_argument('--model-path', type=str, required=True,
help='预训练模型路径')
parser.add_argument('--model-type', type=str, default='fcn8s', choices=['fcn8s', 'fcn16s', 'fcn32s'],
help='FCN模型类型 (fcn8s, fcn16s, fcn32s)')
parser.add_argument('--image-path', type=str, required=True,
help='输入图像路径,可以是单个图像或者目录')
parser.add_argument('--output-dir', type=str, default='results',
help='结果保存目录')
parser.add_argument('--overlay', action='store_true',
help='是否将分割结果与原图叠加')
parser.add_argument('--no-cuda', action='store_true',
help='禁用CUDA')
return parser.parse_args()
3.4.2 图像预处理和后处理函数
接下来,我们定义两个辅助函数:一个用于预处理输入图像,使其符合模型的输入要求;另一个用于将分割结果与原图叠加,增强可视化效果。
python
def preprocess_image(image_path):
"""预处理输入图像"""
image = Image.open(image_path).convert('RGB')
# 图像预处理变换
transform = transforms.Compose([
transforms.Resize(320),
transforms.CenterCrop(320),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
input_tensor = transform(image)
input_batch = input_tensor.unsqueeze(0)
return input_batch, image
def overlay_segmentation(image, segmentation, alpha=0.7):
"""将分割结果与原图叠加"""
image_np = np.array(image)
segmentation_resized = np.array(Image.fromarray(segmentation.astype(np.uint8)).resize(
(image_np.shape[1], image_np.shape[0]), Image.NEAREST))
overlay = image_np.copy()
for i in range(3):
overlay[:, :, i] = image_np[:, :, i] * (1 - alpha) + segmentation_resized[:, :, i] * alpha
return overlay.astype(np.uint8)
preprocess_image
函数将输入图像调整为统一大小(320×320),转换为张量格式,并应用ImageNet数据集的标准归一化。overlay_segmentation
函数则接受原始图像和分割图,按指定的透明度(默认0.7)将它们叠加在一起,使得分割结果更直观。
3.4.3 预测及可视化函数
下面我们实现对图像进行预测和可视化的功能:
python
def predict_image(model, image_path, device, overlay=False):
"""对单个图像进行预测"""
input_batch, original_image = preprocess_image(image_path)
input_batch = input_batch.to(device)
with torch.no_grad():
output = model(input_batch)
output = torch.nn.functional.softmax(output, dim=1)
_, pred = torch.max(output, 1)
pred = pred.cpu().numpy()[0]
# 将预测结果转换为彩色分割图
segmentation_map = decode_segmap(pred)
if overlay:
result = overlay_segmentation(original_image, segmentation_map)
else:
result = segmentation_map
return result, pred, original_image
def predict_and_visualize(model, image_path, output_dir, device, overlay=False):
"""预测图像并可视化结果"""
# 如果图像路径是目录
if os.path.isdir(image_path):
image_files = [f for f in os.listdir(image_path) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
for image_file in image_files:
file_path = os.path.join(image_path, image_file)
visualize_prediction(model, file_path, output_dir, device, overlay)
else:
visualize_prediction(model, image_path, output_dir, device, overlay)
def visualize_prediction(model, image_path, output_dir, device, overlay=False):
"""可视化单个图像的预测结果"""
os.makedirs(output_dir, exist_ok=True)
# 预测图像
result, pred, original_image = predict_image(model, image_path, device, overlay)
# 保存结果
base_name = os.path.basename(image_path).split('.')[0]
plt.figure(figsize=(15, 5))
plt.subplot(1, 3, 1)
plt.title('原始图像')
plt.imshow(original_image)
plt.axis('off')
plt.subplot(1, 3, 2)
plt.title('分割结果')
plt.imshow(decode_segmap(pred))
plt.axis('off')
plt.subplot(1, 3, 3)
if overlay:
plt.title('叠加结果')
else:
plt.title('分割结果')
plt.imshow(result)
plt.axis('off')
plt.tight_layout()
plt.savefig(os.path.join(output_dir, f'{base_name}_result.png'))
plt.close()
class_pixels = {}
for i, class_name in enumerate(VOC_CLASSES):
num_pixels = np.sum(pred == i)
if num_pixels > 0:
class_pixels[class_name] = num_pixels
# 创建类别分布饼图
if class_pixels:
plt.figure(figsize=(10, 10))
labels = list(class_pixels.keys())
sizes = list(class_pixels.values())
plt.pie(sizes, labels=labels, autopct='%1.1f%%', shadow=True, startangle=90)
plt.axis('equal')
plt.title('类别分布')
plt.savefig(os.path.join(output_dir, f'{base_name}_class_dist.png'))
plt.close()
# 保存单独的分割图
segmentation_img = Image.fromarray(decode_segmap(pred))
segmentation_img.save(os.path.join(output_dir, f'{base_name}_segmentation.png'))
预测函数将首先加载并预处理图像,然后通过模型进行前向传播,获取预测结果。而预测结果先通过softmax转换为概率分布,然后选取概率最高的类别作为最终预测。最后,根据是否需要叠加展示,返回相应的可视化结果。
3.4.4 主函数实现
最后,我们实现主函数,将所有功能整合起来,主函数首先解析命令行参数,然后根据参数创建相应的FCN模型。在加载预训练权重时,我特别考虑了PyTorch不同版本的兼容性问题,使用了try-except结构来适应不同版本的加载方式。加载完模型后,将其设置为评估模式,然后调用预测和可视化函数处理指定的图像或图像目录。
python
def main():
args = parse_args()
device = torch.device('cuda' if torch.cuda.is_available() and not args.no_cuda else 'cpu')
print(f'使用设备: {device}')
# 创建模型
model = get_fcn_model(model_type=args.model_type, num_classes=NUM_CLASSES, pretrained=False)
checkpoint = torch.load(args.model_path, map_location=device)
if 'model_state_dict' in checkpoint:
model.load_state_dict(checkpoint['model_state_dict'])
print(f'加载检查点: Epoch {checkpoint["epoch"]}, mIoU {checkpoint["best_miou"]:.4f}')
else:
model.load_state_dict(checkpoint)
print(f'加载模型权重成功')
model = model.to(device)
model.eval()
predict_and_visualize(model, args.image_path, args.output_dir, device, args.overlay)
print(f'结果已保存到: {args.output_dir}')
完成了预测推理代码之后我们便可以使用如下命令进行推理:
bash
python predict.py --model-path checkpoints/fcn8s_best.pth --image-path test_images/test.jpg --overlay

之后我们可以看到我们的目录下面新增了一个results
文件夹用于储存我们的推理结果




可以看到训练了22epoch的效果并不是很理想,目前还只是单纯的训练,没有去深入调优模型超参数和训练策略。实际上,FCN网络的性能还有很大的提升空间。大家可以自己优化一下分割的效果哦