找回密码
 立即注册
首页 业界区 业界 手把手带你从论文出发实战搭建分割FCN网络 ...

手把手带你从论文出发实战搭建分割FCN网络

巫雪艷 4 天前
作者:SkyXZ
CSDN:SkyXZ~-CSDN博客
博客园:SkyXZ - 博客园


  • FCN论文地址:Fully Convolutional Networks for Semantic Segmentation
一、什么是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等更精细的改进版本。
1.png


二、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 1Conv1_1 → Conv1_2 → MaxPool112×112Block 2Conv2_1 → Conv2_2 → MaxPool56×56Block 3Conv3_1 → Conv3_2 → Conv3_3 → MaxPool28×28Block 4Conv4_1 → Conv4_2 → Conv4_3 → MaxPool14×14Block 5Conv5_1 → Conv5_2 → Conv5_3 → MaxPool7×7Classifierfc6(4096) → fc7(4096) → fc8(1000)1×1
注:fc8 的输出通常对应于 ImageNet 的1000个类别。
3.png

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

        这样,网络的输出就变成了一张尺寸更小但仍保留空间结构的 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×2241×Conv1 → Pool1112×1122×Conv2 → Pool256×564×Conv3 → Pool328×288×Conv4 → Pool414×1416×Conv5 → Pool57×732×        也就是说,输入一张 512×512 的图像,经过 VGG16 后输出的特征图仅为 16×16,最终我们得到的语义特征图(即 conv5 输出)只有输入图尺寸的 1/32 大小。意味着每个位置预测的结果实际上对应输入图上的一个 32×32 的区域(感受野即网络在该位置所看到的输入图像区域),这对于物体边缘或细小结构来说是非常粗糙的。由于语义分割的目标是:为图像中的每一个像素分配一个语义标签。下采样过度的网络会导致输出的特征图空间尺寸过小,使得每个像素的预测实际上代表了输入图像上的一个较大区域。如果网络最后输出的特征图过小,我们只能得到非常稀疏的分类结果,哪怕后续再通过插值或上采样恢复图像尺寸,也会因为高频细节已经丢失而无法精确还原边界。
5.png

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

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)详细推导 - 知乎
7.png

2.4 跳跃连接(Skip Connections):细节与语义的结合

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

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

三、实战搭建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
        我们首先配置我们的开发环境:
  1. # step1:创建一个Conda环境并激活
  2. conda create -n FCN python=3.10 -y
  3. conda activate FCN
  4. # step2:下载安装依赖
  5. pip install torch==2.0.1 torchvision==0.15.2
  6. pip install numpy opencv-python matplotlib tqdm pillow
  7. # step3:验证
  8. python -c "import torch; print(torch.__version__, torch.cuda.is_available())"
复制代码
        验证之后如果终端显示True,即代表我们的Pytorch安装正确,如若遇到错误请自行搜索解决方法
10.png

        数据集方面我们选择使用PASCAL VOC 2012,这是语义分割的经典基准数据集,大概有2GB大小,其下载方式如下,同时其解压后的目录格式大致如下:
  1. # step1:下载PASCAL VOC 2012数据集
  2. wget http://host.robots.ox.ac.uk/pascal/VOC/voc2012/VOCtrainval_11-May-2012.tar
  3. # step2:解压
  4. tar -vxf VOCtrainval_11-May-2012.tar
  5. # step3:验证目录结构
  6. 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)等宏变量
  1. import os # 导入os模块,用于文件路径操作
  2. import numpy as np # 导入numpy模块,用于数值计算
  3. import torch # 导入PyTorch主模块
  4. from torch.utils.data import Dataset, DataLoader # 导入数据集和数据加载器
  5. from PIL import Image # 导入PIL图像处理库
  6. import torchvision.transforms as transforms # 导入图像变换模块
  7. from torchvision.transforms import functional as F # 导入函数式变换模块
  8. # VOC数据集的类别名称(21个类别,包括背景)
  9. VOC_CLASSES = [
  10.     'background', 'aeroplane', 'bicycle', 'bird', 'boat', 'bottle',
  11.     'bus', 'car', 'cat', 'chair', 'cow', 'diningtable', 'dog',
  12.     'horse', 'motorbike', 'person', 'potted plant', 'sheep', 'sofa',
  13.     'train', 'tv/monitor'
  14. ]
  15. # 宏定义获取类别数量
  16. NUM_CLASSES = len(VOC_CLASSES)
  17. # 定义PIL的重采样常量
  18. PIL_NEAREST = 0  # 最近邻重采样方式,保持锐利边缘,适用于掩码
  19. PIL_BILINEAR = 1  # 双线性重采样方式,平滑图像,适用于原始图像
  20. # 定义VOC数据集的颜色映射 (用于可视化分割结果),每个类别对应一个RGB颜色
  21. VOC_COLORMAP = [
  22.     [0, 0, 0], [128, 0, 0], [0, 128, 0], [128, 128, 0], [0, 0, 128],
  23.     [128, 0, 128], [0, 128, 128], [128, 128, 128], [64, 0, 0], [192, 0, 0],
  24.     [64, 128, 0], [192, 128, 0], [64, 0, 128], [192, 0, 128], [64, 128, 128],
  25.     [192, 128, 128], [0, 64, 0], [128, 64, 0], [0, 192, 0], [128, 192, 0],
  26.     [0, 64, 128]
  27. ]
复制代码
3.1.2 创建VOCSegmentation(Dataset)类

        接着我们创建一个类 VOCSegmentation(Dataset),用于封装和加载 VOC2012 语义分割数据集。该类继承自 PyTorch 的 Dataset,实现了标准的 __getitem__ 和 __len__ 方法,可直接配合 DataLoader 批量加载数据。它能够根据数据划分文件(如 train.txt、val.txt)读取图像与对应的掩码路径,并对它们进行统一尺寸的预处理——图像采用双线性插值以保持平滑性,掩码则使用最近邻插值以避免引入伪标签。同时,彩色掩码图会根据预定义的 VOC_COLORMAP 进行颜色到类别索引的映射,最终转换为模型训练所需的二维整数张量,实现从 RGB 掩码到语义标签的准确转换。
  1. class VOCSegmentation(Dataset):
  2.     """
  3.     VOC2012语义分割数据集的PyTorch Dataset实现
  4.     负责数据的加载、预处理和转换
  5.     """
  6.     def __init__(self, root, split='train', transform=None, target_transform=None, img_size=320):
  7.         """
  8.         初始化数据集
  9.         参数:
  10.             root (string): VOC数据集的根目录路径
  11.             split (string, optional): 使用的数据集划分,可选 'train', 'val' 或 'trainval'
  12.             transform (callable, optional): 输入图像的变换函数
  13.             target_transform (callable, optional): 目标掩码的变换函数
  14.             img_size (int, optional): 调整图像和掩码的大小
  15.         """
  16.     def __getitem__(self, index):
  17.         """
  18.         获取数据集中的一个样本
  19.         参数:
  20.             index (int): 样本索引
  21.         返回:
  22.             tuple: (图像, 掩码) 对,分别为图像张量和掩码张量
  23.         """
  24.     def __len__(self):
  25.         """返回数据集中的样本数量"""
复制代码
        我们来具体实现这三个类函数,首先是 __init__。在该构造函数中,我们传入数据集的根目录 root、划分类型 split(如 'train'、'val' 或 'trainval')、图像变换 transform、标签变换 target_transform 以及目标尺寸 img_size。随后根据划分文件(如 train.txt)读取图像和掩码的文件名,拼接得到完整路径,并对路径有效性进行检查。最后,我们将图像和掩码路径分别保存在 self.images 和 self.masks 中,便于后续索引使用。
  1. def __init__(self, root, split='train', transform=None, target_transform=None, img_size=320):
  2.         super(VOCSegmentation, self).__init__()
  3.         self.root = root
  4.         self.split = split
  5.         self.transform = transform
  6.         self.target_transform = target_transform
  7.         self.img_size = img_size
  8.         # 确定图像和标签文件的路径
  9.         voc_root = self.root
  10.         image_dir = os.path.join(voc_root, 'JPEGImages')  # 原始图像目录
  11.         mask_dir = os.path.join(voc_root, 'SegmentationClass')  # 语义分割标注目录
  12.         # 获取图像文件名列表(从划分文件中读取)
  13.         splits_dir = os.path.join(voc_root, 'ImageSets', 'Segmentation')
  14.         split_file = os.path.join(splits_dir, self.split + '.txt')
  15.         # 确保分割文件存在
  16.         if not os.path.exists(split_file):
  17.             raise FileNotFoundError(f"找不到拆分文件: {split_file}")
  18.         # 读取文件名列表
  19.         with open(split_file, 'r') as f:
  20.             file_names = [x.strip() for x in f.readlines()]
  21.         # 构建图像和掩码的完整路径
  22.         self.images = [os.path.join(image_dir, x + '.jpg') for x in file_names]
  23.         self.masks = [os.path.join(mask_dir, x + '.png') for x in file_names]
  24.         # 检查文件是否存在,打印警告但不中断程序
  25.         for img_path in self.images:
  26.             if not os.path.exists(img_path):
  27.                 print(f"警告: 图像文件不存在: {img_path}")
  28.         for mask_path in self.masks:
  29.             if not os.path.exists(mask_path):
  30.                 print(f"警告: 掩码文件不存在: {mask_path}")
  31.         # 确保图像和掩码数量匹配
  32.         assert len(self.images) == len(self.masks), "图像和掩码数量不匹配"
  33.         print(f"加载了 {len(self.images)} 对图像和掩码用于{split}集")
复制代码
        接下来是 __getitem__ 方法,用于根据索引加载一个样本。图像和掩码被读取并转换为 RGB 格式,再统一调整为设定大小,其中图像使用双线性插值以保持平滑性,掩码使用最近邻插值以保留类别标签。图像经过指定的变换函数处理后,掩码则根据是否提供 target_transform 进行处理;若未指定,我们将掩码由 RGB 图转为类别索引图,通过遍历预定义的 VOC_COLORMAP 映射每个像素所属的语义类别,最终转为 long 类型的 PyTorch 张量,便于模型训练使用。最后的__len__ 方法比较简单,直接返回数据集中图像的总数就好了,也就是 self.images 的长度,用于告知 PyTorch DataLoader 数据集的大小。
  1. def __getitem__(self, index):
  2.         # 加载图像和掩码
  3.         img_path = self.images[index]
  4.         mask_path = self.masks[index]
  5.         img = Image.open(img_path).convert('RGB')# 打开图像并转换为RGB格式
  6.         mask = Image.open(mask_path).convert('RGB')# 打开掩码并转换为RGB格式(确保与colormap匹配)
  7.         # 统一调整图像和掩码大小,确保尺寸一致
  8.         img = img.resize((self.img_size, self.img_size), PIL_BILINEAR)# 对于图像使用双线性插值以保持平滑
  9.         mask = mask.resize((self.img_size, self.img_size), PIL_NEAREST)# 对于掩码使用最近邻插值以避免引入新的类别值
  10.         # 应用图像变换
  11.         if self.transform is not None:
  12.             img = self.transform(img)
  13.         # 处理掩码变换
  14.         if self.target_transform is not None:
  15.             mask = self.target_transform(mask)
  16.         else:
  17.             # 将掩码转换为类别索引
  18.             mask = np.array(mask)
  19.             # 检查掩码的维度,确保是RGB(3通道)
  20.             if len(mask.shape) != 3 or mask.shape[2] != 3:
  21.                 raise ValueError(f"掩码维度错误: {mask.shape}, 期望为 (H,W,3)")
  22.             mask_copy = np.zeros((mask.shape[0], mask.shape[1]), dtype=np.uint8)# 创建一个新的类别索引掩码
  23.             # 将RGB颜色映射到类别索引
  24.             # 遍历每种颜色,将对应像素设置为类别索引
  25.             for k, color in enumerate(VOC_COLORMAP):
  26.                 # 将每个颜色通道转换为布尔掩码
  27.                 r_match = mask[:, :, 0] == color[0]
  28.                 g_match = mask[:, :, 1] == color[1]
  29.                 b_match = mask[:, :, 2] == color[2]
  30.                 # 只有三个通道都匹配的像素才被分配为此类别
  31.                 color_match = r_match & g_match & b_match
  32.                 mask_copy[color_match] = k
  33.             mask = torch.from_numpy(mask_copy).long() # 转换为PyTorch张量(长整型,用于交叉熵损失)
  34.         return img, mask
  35.     def __len__(self):
  36.         return len(self.images)
复制代码
因此完整的类代码如下:
  1. class VOCSegmentation(Dataset):    def __init__(self, root, split='train', transform=None, target_transform=None, img_size=320):
  2.         super(VOCSegmentation, self).__init__()
  3.         self.root = root
  4.         self.split = split
  5.         self.transform = transform
  6.         self.target_transform = target_transform
  7.         self.img_size = img_size
  8.         # 确定图像和标签文件的路径
  9.         voc_root = self.root
  10.         image_dir = os.path.join(voc_root, 'JPEGImages')  # 原始图像目录
  11.         mask_dir = os.path.join(voc_root, 'SegmentationClass')  # 语义分割标注目录
  12.         # 获取图像文件名列表(从划分文件中读取)
  13.         splits_dir = os.path.join(voc_root, 'ImageSets', 'Segmentation')
  14.         split_file = os.path.join(splits_dir, self.split + '.txt')
  15.         # 确保分割文件存在
  16.         if not os.path.exists(split_file):
  17.             raise FileNotFoundError(f"找不到拆分文件: {split_file}")
  18.         # 读取文件名列表
  19.         with open(split_file, 'r') as f:
  20.             file_names = [x.strip() for x in f.readlines()]
  21.         # 构建图像和掩码的完整路径
  22.         self.images = [os.path.join(image_dir, x + '.jpg') for x in file_names]
  23.         self.masks = [os.path.join(mask_dir, x + '.png') for x in file_names]
  24.         # 检查文件是否存在,打印警告但不中断程序
  25.         for img_path in self.images:
  26.             if not os.path.exists(img_path):
  27.                 print(f"警告: 图像文件不存在: {img_path}")
  28.         for mask_path in self.masks:
  29.             if not os.path.exists(mask_path):
  30.                 print(f"警告: 掩码文件不存在: {mask_path}")
  31.         # 确保图像和掩码数量匹配
  32.         assert len(self.images) == len(self.masks), "图像和掩码数量不匹配"
  33.         print(f"加载了 {len(self.images)} 对图像和掩码用于{split}集")    def __getitem__(self, index):
  34.         # 加载图像和掩码
  35.         img_path = self.images[index]
  36.         mask_path = self.masks[index]
  37.         img = Image.open(img_path).convert('RGB')# 打开图像并转换为RGB格式
  38.         mask = Image.open(mask_path).convert('RGB')# 打开掩码并转换为RGB格式(确保与colormap匹配)
  39.         # 统一调整图像和掩码大小,确保尺寸一致
  40.         img = img.resize((self.img_size, self.img_size), PIL_BILINEAR)# 对于图像使用双线性插值以保持平滑
  41.         mask = mask.resize((self.img_size, self.img_size), PIL_NEAREST)# 对于掩码使用最近邻插值以避免引入新的类别值
  42.         # 应用图像变换
  43.         if self.transform is not None:
  44.             img = self.transform(img)
  45.         # 处理掩码变换
  46.         if self.target_transform is not None:
  47.             mask = self.target_transform(mask)
  48.         else:
  49.             # 将掩码转换为类别索引
  50.             mask = np.array(mask)
  51.             # 检查掩码的维度,确保是RGB(3通道)
  52.             if len(mask.shape) != 3 or mask.shape[2] != 3:
  53.                 raise ValueError(f"掩码维度错误: {mask.shape}, 期望为 (H,W,3)")
  54.             mask_copy = np.zeros((mask.shape[0], mask.shape[1]), dtype=np.uint8)# 创建一个新的类别索引掩码
  55.             # 将RGB颜色映射到类别索引
  56.             # 遍历每种颜色,将对应像素设置为类别索引
  57.             for k, color in enumerate(VOC_COLORMAP):
  58.                 # 将每个颜色通道转换为布尔掩码
  59.                 r_match = mask[:, :, 0] == color[0]
  60.                 g_match = mask[:, :, 1] == color[1]
  61.                 b_match = mask[:, :, 2] == color[2]
  62.                 # 只有三个通道都匹配的像素才被分配为此类别
  63.                 color_match = r_match & g_match & b_match
  64.                 mask_copy[color_match] = k
  65.             mask = torch.from_numpy(mask_copy).long() # 转换为PyTorch张量(长整型,用于交叉熵损失)
  66.         return img, mask
  67.     def __len__(self):
  68.         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 即可。
  1. def get_transforms(train=True):
  2.     """
  3.     获取图像变换函数
  4.     参数:
  5.         train (bool): 是否为训练集,决定是否应用数据增强
  6.     返回:
  7.         tuple: (图像变换, 目标掩码变换)
  8.     """
  9.     if train:
  10.         # 训练集使用数据增强
  11.         transform = transforms.Compose([
  12.             transforms.RandomHorizontalFlip(),# 随机水平翻转增加数据多样性
  13.             transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),# 随机调整亮度、对比度和饱和度
  14.             transforms.ToTensor(),# 转换为张量(值范围变为[0,1])
  15.             transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) # 归一化(使用ImageNet的均值和标准差)
  16.         ])
  17.     else:
  18.         # 验证集只需要基本变换
  19.         transform = transforms.Compose([
  20.             transforms.ToTensor(), # 转换为张量
  21.             transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])# 归一化
  22.         ])
  23.     target_transform = None# 掩码不需要标准化或者转换为Tensor (已在__getitem__中处理)
  24.     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 的数据拷贝效率。
        该函数最终返回构建好的训练与验证加载器,可直接用于训练循环中的迭代操作,完整代码如下所示:
  1. def get_data_loaders(voc_root, batch_size=4, num_workers=4, img_size=320):
  2.     """
  3.     创建训练和验证数据加载器
  4.     参数:
  5.         voc_root (string): VOC数据集根目录
  6.         batch_size (int): 批次大小
  7.         num_workers (int): 数据加载的线程数
  8.         img_size (int): 图像的大小
  9.     返回:
  10.         tuple: (train_loader, val_loader) 训练和验证数据加载器
  11.     """
  12.     # 获取图像和掩码变换
  13.     train_transform, train_target_transform = get_transforms(train=True)
  14.     val_transform, val_target_transform = get_transforms(train=False)
  15.     # 创建训练数据集
  16.     train_dataset = VOCSegmentation(
  17.         root=voc_root,
  18.         split='train',  # 使用训练集划分
  19.         transform=train_transform,
  20.         target_transform=train_target_transform,
  21.         img_size=img_size
  22.     )
  23.     # 创建验证数据集
  24.     val_dataset = VOCSegmentation(
  25.         root=voc_root,
  26.         split='val',  # 使用验证集划分
  27.         transform=val_transform,
  28.         target_transform=val_target_transform,
  29.         img_size=img_size
  30.     )
  31.     # 创建训练数据加载器
  32.     train_loader = DataLoader(
  33.         train_dataset,
  34.         batch_size=batch_size,
  35.         shuffle=True,  # 随机打乱数据
  36.         num_workers=num_workers,  # 多线程加载
  37.         pin_memory=True,  # 数据预加载到固定内存,加速GPU传输
  38.         drop_last=True  # 丢弃最后不足一个批次的数据
  39.     )
  40.     # 创建验证数据加载器
  41.     val_loader = DataLoader(
  42.         val_dataset,
  43.         batch_size=batch_size,
  44.         shuffle=False,  # 不打乱数据
  45.         num_workers=num_workers,
  46.         pin_memory=True
  47.     )
  48.     return train_loader, val_loader
复制代码
3.1.5 可视化分割结果

        在训练语义分割模型的过程中,如果我们只是单纯地关注每轮训练的数值指标(如 IoU、准确率等),难免会显得有些枯燥,且难以直观感受模型到底学得怎么样。尤其是在模型逐步收敛时,仅靠指标的波动并不能很好地揭示模型的细节表现。因此,我在此基础上引入了一个可视化辅助函数 decode_segmap,用于将模型预测得到的分割结果从类别索引图转换为彩色图像。这样一来,我们就可以将每个像素所属的类别清晰地呈现在图像上,借助这个工具,我们可以在训练过程中插入实时可视化,随时查看模型对于不同样本的分割表现,为调参和模型改进提供更加直观的反馈。完整实现的代码如下:
  1. def decode_segmap(segmap):
  2.     """
  3.     将类别索引的分割图转换为RGB彩色图像(用于可视化)
  4.     参数:
  5.         segmap (np.array或torch.Tensor): 形状为(H,W)的分割图,值为类别索引
  6.     返回:
  7.         rgb_img (np.array): 形状为(H,W,3)的RGB彩色图像
  8.     """
  9.     # 确保segmap是NumPy数组
  10.     if isinstance(segmap, torch.Tensor):
  11.         segmap = segmap.cpu().numpy()
  12.     # 检查segmap的形状,处理各种可能的输入格式
  13.     if len(segmap.shape) > 2:
  14.         if len(segmap.shape) == 3 and segmap.shape[0] <= 3:
  15.             segmap = segmap[0]
  16.     rgb_img = np.zeros((segmap.shape[0], segmap.shape[1], 3), dtype=np.uint8)    # 创建RGB图像
  17.     # 根据类别索引填充对应的颜色
  18.     for cls_idx, color in enumerate(VOC_COLORMAP):
  19.         mask = segmap == cls_idx# 找到属于当前类别的像素
  20.         if mask.any():  # 只处理存在的类别
  21.             rgb_img[mask] = color # 将这些像素设置为对应的颜色
  22.     return rgb_img
复制代码
3.3.3 分割结果可视化函数

        为了直观地观察分割结果,我们定义一个函数用于保存预测结果的可视化图像。这个函数将从验证集中取出一定数量的样本,通过模型进行预测,然后将原始图像、真实标签和预测结果并排显示并保存为图像文件,这样我们就可以直观地观察模型的分割效果
  1. import torch
  2. import torch.nn as nn
  3. import torch.nn.functional as F
  4. import torchvision.models as models
  5. from dataload import NUM_CLASSES
复制代码
3.3.4 主训练函数

        最后,我们编写主函数,实现完整的训练流程。主函数是整个训练脚本的核心,它将各个组件有机地整合在一起,形成完整的训练流程。首先,它通过parse_args()解析命令行输入的各项参数,如数据集路径、模型类型、批量大小等,使训练过程更加灵活可控。之后,它会调用get_data_loaders()函数加载并预处理VOC数据集,同时创建数据加载器以便批量获取训练和验证样本。接着,根据参入参数指定的模型类型(FCN8s/FCN16s/FCN32s)实例化相应的网络结构,并将其转移到可用的计算设备(GPU或CPU)上。在优化策略方面,主函数使用交叉熵损失函数(忽略255标签值)评估分割质量,采用带动量的SGD优化器更新网络参数,并通过学习率调度器在训练后期降低学习率以获得更精细的优化效果。如果参入参数提供了检查点路径,函数会从中恢复模型权重、优化器状态和训练进度,实现断点续训。核心的训练循环涵盖了完整的训练-评估-保存流程:每个epoch内先在训练集上进行前向传播、损失计算、反向传播和参数更新;然后在验证集上评估模型性能(损失值、像素准确率和mIoU);当取得更高mIoU时,保存最佳模型并生成可视化结果,同时定期保存最新模型以防训练中断。训练完成后,主函数会绘制整个训练过程的损失曲线、准确率曲线和mIoU曲线,直观展示模型的学习轨迹和性能变化,帮助大家更好地理解训练动态并优化训练策略。
  1. def _initialize_weights(self):
  2.         # 初始化反卷积层的权重为双线性上采样
  3.         for m in self.modules():
  4.             if isinstance(m, nn.ConvTranspose2d):
  5.                 # 双线性上采样的初始化
  6.                 m.weight.data.zero_()
  7.                 m.weight.data = self._make_bilinear_weights(m.kernel_size[0], m.out_channels)
复制代码
        完成了训练代码之后我们便可以开始训练啦!我们输入以下命令即可开始训练:
  1. def _make_bilinear_weights(self, size, num_channels):
  2.         """生成双线性插值的权重"""
  3.         factor = (size + 1) // 2
  4.         if size % 2 == 1:
  5.             center = factor - 1
  6.         else:
  7.             center = factor - 0.5
  8.         og = torch.FloatTensor(size, size)
  9.         for i in range(size):
  10.             for j in range(size):
  11.                 og[i, j] = (1 - abs((i - center) / factor)) * (1 - abs((j - center) / factor))
  12.         filter = torch.zeros(num_channels, num_channels, size, size)
  13.         for i in range(num_channels):
  14.             filter[i, i] = og
  15.         return filter
复制代码
11.png

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

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

14.png

15.png

3.4 完成推理预测脚本predict.py

        训练好模型后,我们需要一个单独的脚本来对新图像进行语义分割预测。这个推理脚本不仅能够加载我们训练好的模型,还能对单张图像或整个文件夹的图像进行批量预测,同时提供多种可视化方式展示分割结果。下面我将详细讲解这个推理脚本的实现过程。
3.4.1 导入必要模块和解析命令行参数

        首先,我们需要导入必要的模块,并设置命令行参数解析器,以便灵活地配置推理过程。在参数解析部分,我们可以通过--model-path指定预训练模型的存储路径;同时通过--model-type选择使用FCN8s、FCN16s或FCN32s中的任一模型架构。而--image-path参数支持单个图像文件也可以制定一个文件夹进行批量处理。分割结果默认保存在名为"results"的文件夹中,也可以通过--output-dir参数自定义存储位置,--overlay参数则可以选择是否将掩码叠加在原图上面
  1. class FCN32s(nn.Module):
  2.     def __init__(self, num_classes=NUM_CLASSES, pretrained=True):
  3.         super(FCN32s, self).__init__()
复制代码
3.4.2 图像预处理和后处理函数

        接下来,我们定义两个辅助函数:一个用于预处理输入图像,使其符合模型的输入要求;另一个用于将分割结果与原图叠加,增强可视化效果。
  1.                 vgg16 = models.vgg16(pretrained=pretrained)# 加载预训练的VGG16模型
  2.         features = list(vgg16.features.children())# 获取特征提取部分
  3.         # 根据FCN原始论文修改VGG16网络
  4.         # 前5段卷积块保持不变
  5.         self.features1 = nn.Sequential(*features[:5])    # conv1 + pool1
  6.         self.features2 = nn.Sequential(*features[5:10])  # conv2 + pool2
  7.         self.features3 = nn.Sequential(*features[10:17]) # conv3 + pool3
  8.         self.features4 = nn.Sequential(*features[17:24]) # conv4 + pool4
  9.         self.features5 = nn.Sequential(*features[24:31]) # conv5 + pool5
  10.         # 全连接层替换为1x1卷积
  11.         self.fc6 = nn.Conv2d(512, 4096, kernel_size=7, padding=3)
  12.         self.relu6 = nn.ReLU(inplace=True)
  13.         self.drop6 = nn.Dropout2d()
  14.         self.fc7 = nn.Conv2d(4096, 4096, kernel_size=1)
  15.         self.relu7 = nn.ReLU(inplace=True)
  16.         self.drop7 = nn.Dropout2d()
  17.         # 分类层
  18.         self.score = nn.Conv2d(4096, num_classes, kernel_size=1)
  19.         # 上采样层: 32倍上采样回原始图像大小
  20.         self.upsample = nn.ConvTranspose2d(num_classes, num_classes, kernel_size=64, stride=32, padding=16, bias=False)
  21.         # 初始化参数
  22.         self._initialize_weights()
复制代码
        preprocess_image函数将输入图像调整为统一大小(320×320),转换为张量格式,并应用ImageNet数据集的标准归一化。overlay_segmentation函数则接受原始图像和分割图,按指定的透明度(默认0.7)将它们叠加在一起,使得分割结果更直观。
3.4.3 预测及可视化函数

        下面我们实现对图像进行预测和可视化的功能:
  1. def forward(self, x):
  2.     input_size = x.size()[2:]# 记录输入尺寸用于上采样
  3.     # 编码器 (VGG16)
  4.     x = self.features1(x)
  5.     x = self.features2(x)
  6.     x = self.features3(x)
  7.     x = self.features4(x)
  8.     x = self.features5(x)
  9.     # 全连接层 (以卷积形式实现)
  10.     x = self.relu6(self.fc6(x))
  11.     x = self.drop6(x)
  12.     x = self.relu7(self.fc7(x))
  13.     x = self.drop7(x)
  14.     x = self.score(x)# 分类
  15.     x = self.upsample(x)# 上采样回原始尺寸
  16.     x = x[:, :, :input_size[0], :input_size[1]]# 裁剪到原始图像尺寸
  17.     return x
复制代码
        预测函数将首先加载并预处理图像,然后通过模型进行前向传播,获取预测结果。而预测结果先通过softmax转换为概率分布,然后选取概率最高的类别作为最终预测。最后,根据是否需要叠加展示,返回相应的可视化结果。
3.4.4 主函数实现

        最后,我们实现主函数,将所有功能整合起来,主函数首先解析命令行参数,然后根据参数创建相应的FCN模型。在加载预训练权重时,我特别考虑了PyTorch不同版本的兼容性问题,使用了try-except结构来适应不同版本的加载方式。加载完模型后,将其设置为评估模式,然后调用预测和可视化函数处理指定的图像或图像目录。
  1. class FCN32s(nn.Module):
  2.     def __init__(self, num_classes=NUM_CLASSES, pretrained=True):
  3.         super(FCN32s, self).__init__()
  4.         vgg16 = models.vgg16(pretrained=pretrained)# 加载预训练的VGG16模型
  5.         features = list(vgg16.features.children())# 获取特征提取部分
  6.         # 根据FCN原始论文修改VGG16网络
  7.         # 前5段卷积块保持不变
  8.         self.features1 = nn.Sequential(*features[:5])    # conv1 + pool1
  9.         self.features2 = nn.Sequential(*features[5:10])  # conv2 + pool2
  10.         self.features3 = nn.Sequential(*features[10:17]) # conv3 + pool3
  11.         self.features4 = nn.Sequential(*features[17:24]) # conv4 + pool4
  12.         self.features5 = nn.Sequential(*features[24:31]) # conv5 + pool5
  13.         # 全连接层替换为1x1卷积
  14.         self.fc6 = nn.Conv2d(512, 4096, kernel_size=7, padding=3)
  15.         self.relu6 = nn.ReLU(inplace=True)
  16.         self.drop6 = nn.Dropout2d()
  17.         self.fc7 = nn.Conv2d(4096, 4096, kernel_size=1)
  18.         self.relu7 = nn.ReLU(inplace=True)
  19.         self.drop7 = nn.Dropout2d()
  20.         self.score = nn.Conv2d(4096, num_classes, kernel_size=1)# 分类层
  21.         # 上采样层: 32倍上采样回原始图像大小
  22.         self.upsample = nn.ConvTranspose2d(num_classes, num_classes, kernel_size=64, stride=32, padding=16, bias=False)
  23.         self._initialize_weights()# 初始化参数
  24.     def forward(self, x):
  25.         input_size = x.size()[2:]# 记录输入尺寸用于上采样
  26.         # 编码器 (VGG16)
  27.         x = self.features1(x)
  28.         x = self.features2(x)
  29.         x = self.features3(x)
  30.         x = self.features4(x)
  31.         x = self.features5(x)
  32.         # 全连接层 (以卷积形式实现)
  33.         x = self.relu6(self.fc6(x))
  34.         x = self.drop6(x)
  35.         x = self.relu7(self.fc7(x))
  36.         x = self.drop7(x)
  37.         x = self.score(x)# 分类
  38.         x = self.upsample(x)# 上采样回原始尺寸
  39.         x = x[:, :, :input_size[0], :input_size[1]]# 裁剪到原始图像尺寸
  40.         return x
  41.   
  42.     def _initialize_weights(self):
  43.         # 初始化反卷积层的权重为双线性上采样
  44.         for m in self.modules():
  45.             if isinstance(m, nn.ConvTranspose2d):
  46.                 # 双线性上采样的初始化
  47.                 m.weight.data.zero_()
  48.                 m.weight.data = self._make_bilinear_weights(m.kernel_size[0], m.out_channels)
  49.     def _make_bilinear_weights(self, size, num_channels):
  50.         """生成双线性插值的权重"""
  51.         factor = (size + 1) // 2
  52.         if size % 2 == 1:
  53.             center = factor - 1
  54.         else:
  55.             center = factor - 0.5
  56.         og = torch.FloatTensor(size, size)
  57.         for i in range(size):
  58.             for j in range(size):
  59.                 og[i, j] = (1 - abs((i - center) / factor)) * (1 - abs((j - center) / factor))
  60.         filter = torch.zeros(num_channels, num_channels, size, size)
  61.         for i in range(num_channels):
  62.             filter[i, i] = og
  63.         return filter
复制代码
        完成了预测推理代码之后我们便可以使用如下命令进行推理:
  1. ### FCN16s
  2. class FCN16s(nn.Module):
  3.     def __init__(self, num_classes=NUM_CLASSES, pretrained=True):
  4.         # 获取特征提取部分
  5.         # 分段处理VGG16特征
  6.         ######以上和FCN32s保持一致#########
  7.         # 全连接层替换为1x1卷积
  8.         self.fc6 = nn.Conv2d(512, 4096, kernel_size=7, padding=3)
  9.         self.relu6 = nn.ReLU(inplace=True)
  10.         self.drop6 = nn.Dropout2d()
  11.         self.fc7 = nn.Conv2d(4096, 4096, kernel_size=1)
  12.         self.relu7 = nn.ReLU(inplace=True)
  13.         self.drop7 = nn.Dropout2d()
  14.         self.score_fr = nn.Conv2d(4096, num_classes, kernel_size=1)# 分类层
  15.         self.score_pool4 = nn.Conv2d(512, num_classes, kernel_size=1)# pool4的1x1卷积,用于特征融合
  16.         # 2倍上采样conv7特征
  17.         self.upsample2 = nn.ConvTranspose2d(num_classes, num_classes, kernel_size=4, stride=2, padding=1, bias=False)
  18.         # 16倍上采样回原始图像大小
  19.         self.upsample16 = nn.ConvTranspose2d(num_classes, num_classes, kernel_size=32, stride=16, padding=8, bias=False)
  20.         # 初始化参数
  21.         self._initialize_weights()
  22.     def forward(self, x):
  23.         input_size = x.size()[2:]# 记录输入尺寸用于上采样
  24.         # 编码器 (VGG16)
  25.         x = self.features1(x)
  26.         x = self.features2(x)
  27.         x = self.features3(x)
  28.         pool4 = self.features4(x)# 保存pool4的输出用于后续融合
  29.         x = self.features5(pool4)
  30.         # 全连接层 (以卷积形式实现)
  31.         x = self.relu6(self.fc6(x))
  32.         x = self.drop6(x)
  33.         x = self.relu7(self.fc7(x))
  34.         x = self.drop7(x)
  35.         x = self.score_fr(x)# 分类
  36.         # 2倍上采样
  37.         x = self.upsample2(x)
  38.         # 获取pool4的分数并裁剪
  39.         score_pool4 = self.score_pool4(pool4)
  40.         score_pool4 = score_pool4[:, :, :x.size()[2], :x.size()[3]]
  41.         x = x + score_pool4# 融合特征
  42.         x = self.upsample16(x)# 16倍上采样回原始尺寸
  43.         x = x[:, :, :input_size[0], :input_size[1]]# 裁剪到原始图像尺寸
  44.         return x
  45. ### FCN8s
  46. class FCN8s(nn.Module):
  47.     def __init__(self, num_classes=NUM_CLASSES, pretrained=True):
  48.         # 获取特征提取部分
  49.         # 分段处理VGG16特征
  50.         ######以上和FCN32s保持一致#########
  51.          # 全连接层替换为1x1卷积
  52.         self.fc6 = nn.Conv2d(512, 4096, kernel_size=7, padding=3)
  53.         self.relu6 = nn.ReLU(inplace=True)
  54.         self.drop6 = nn.Dropout2d()
  55.         self.fc7 = nn.Conv2d(4096, 4096, kernel_size=1)
  56.         self.relu7 = nn.ReLU(inplace=True)
  57.         self.drop7 = nn.Dropout2d()
  58.         self.score_fr = nn.Conv2d(4096, num_classes, kernel_size=1)# 分类层
  59.         # pool3和pool4的1x1卷积,用于特征融合
  60.         self.score_pool4 = nn.Conv2d(512, num_classes, kernel_size=1)
  61.         self.score_pool3 = nn.Conv2d(256, num_classes, kernel_size=1)
  62.         # 2倍上采样conv7特征
  63.         self.upsample2_1 = nn.ConvTranspose2d(num_classes, num_classes, kernel_size=4, stride=2, padding=1, bias=False)
  64.         # 2倍上采样融合后的特征
  65.         self.upsample2_2 = nn.ConvTranspose2d(num_classes, num_classes, kernel_size=4, stride=2, padding=1, bias=False)
  66.         # 8倍上采样回原始图像大小
  67.         self.upsample8 = nn.ConvTranspose2d(num_classes, num_classes, kernel_size=16, stride=8, padding=4, bias=False)
  68.         # 初始化参数
  69.         self._initialize_weights()
  70.     def forward(self, x):
  71.         input_size = x.size()[2:]# 记录输入尺寸用于上采样
  72.         # 编码器 (VGG16)
  73.         x = self.features1(x)
  74.         x = self.features2(x)
  75.         pool3 = self.features3(x)# 保存pool3的输出用于后续融合
  76.         pool4 = self.features4(pool3)# 保存pool4的输出用于后续融合
  77.         x = self.features5(pool4)
  78.         # 全连接层 (以卷积形式实现)
  79.         x = self.relu6(self.fc6(x))
  80.         x = self.drop6(x)
  81.         x = self.relu7(self.fc7(x))
  82.         x = self.drop7(x)
  83.         x = self.score_fr(x)# 分类
  84.         x = self.upsample2_1(x)# 2倍上采样
  85.         # 获取pool4的分数并裁剪
  86.         score_pool4 = self.score_pool4(pool4)
  87.         score_pool4 = score_pool4[:, :, :x.size()[2], :x.size()[3]]
  88.         x = x + score_pool4 # 第一次融合特征 (pool5上采样 + pool4)
  89.         x = self.upsample2_2(x)# 再次2倍上采样
  90.         # 获取pool3的分数并裁剪
  91.         score_pool3 = self.score_pool3(pool3)
  92.         score_pool3 = score_pool3[:, :, :x.size()[2], :x.size()[3]]
  93.         x = x + score_pool3# 第二次融合特征 (第一次融合的上采样 + pool3)
  94.         x = self.upsample8(x)# 8倍上采样回原始尺寸
  95.         x = x[:, :, :input_size[0], :input_size[1]]# 裁剪到原始图像尺寸
  96.         return x
复制代码
16.png

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

18.png

19.png

20.png

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

来源:程序园用户自行投稿发布,如果侵权,请联系站长删除
免责声明:如果侵犯了您的权益,请联系站长,我们会及时删除侵权内容,谢谢合作!
您需要登录后才可以回帖 登录 | 立即注册