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

作者：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等更精致的改进版本。
  

二、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个类别。
  

        正是由于全连接层（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类）

        这样，网络的输出就变成了一张尺寸更小但仍保留空间布局的 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 的地区（感受野即网络在该位置所看到的输入图像地区），这对于物体边缘或细小布局来说是非常粗糙的。由于语义分割的目的是：为图像中的每一个像素分配一个语义标签。下采样过分的网络会导致输出的特性图空间尺寸过小，使得每个像素的预测实际上代表了输入图像上的一个较大地区。假如网络最后输出的特性图过小，我们只能得到非常稀疏的分类效果，哪怕后续再通过插值或上采样规复图像尺寸，也会由于高频细节已经丢失而无法正确还原边界。

        为了补充下采样带来的空间精度损失，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
          我们首先配置我们的开发情况:

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安装正确,如若遇到错误请自行搜索解决方法

        数据集方面我们选择使用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 # 导入函数式变换模块
9. # VOC数据集的类别名称（21个类别，包括背景）
10. VOC_CLASSES = [
11.     'background', 'aeroplane', 'bicycle', 'bird', 'boat', 'bottle',
12.     'bus', 'car', 'cat', 'chair', 'cow', 'diningtable', 'dog',
13.     'horse', 'motorbike', 'person', 'potted plant', 'sheep', 'sofa',
14.     'train', 'tv/monitor'
15. ]
16. # 宏定义获取类别数量
17. NUM_CLASSES = len(VOC_CLASSES)
18. # 定义PIL的重采样常量
19. PIL_NEAREST = 0  # 最近邻重采样方式，保持锐利边缘，适用于掩码
20. PIL_BILINEAR = 1  # 双线性重采样方式，平滑图像，适用于原始图像
21. # 定义VOC数据集的颜色映射 (用于可视化分割结果)，每个类别对应一个RGB颜色
22. VOC_COLORMAP = [
23.     [0, 0, 0], [128, 0, 0], [0, 128, 0], [128, 128, 0], [0, 0, 128],
24.     [128, 0, 128], [0, 128, 128], [128, 128, 128], [64, 0, 0], [192, 0, 0],
25.     [64, 128, 0], [192, 128, 0], [64, 0, 128], [192, 0, 128], [64, 128, 128],
26.     [192, 128, 128], [0, 64, 0], [128, 64, 0], [0, 192, 0], [128, 192, 0],
27.     [0, 64, 128]
28. ]

复制代码

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}集")
34.     def __getitem__(self, index):
35.         # 加载图像和掩码
36.         img_path = self.images[index]
37.         mask_path = self.masks[index]
38.         img = Image.open(img_path).convert('RGB')# 打开图像并转换为RGB格式
39.         mask = Image.open(mask_path).convert('RGB')# 打开掩码并转换为RGB格式（确保与colormap匹配）
40.         # 统一调整图像和掩码大小，确保尺寸一致
41.         img = img.resize((self.img_size, self.img_size), PIL_BILINEAR)# 对于图像使用双线性插值以保持平滑
42.         mask = mask.resize((self.img_size, self.img_size), PIL_NEAREST)# 对于掩码使用最近邻插值以避免引入新的类别值
43.         # 应用图像变换
44.         if self.transform is not None:
45.             img = self.transform(img)
46.         # 处理掩码变换
47.         if self.target_transform is not None:
48.             mask = self.target_transform(mask)
49.         else:
50.             # 将掩码转换为类别索引
51.             mask = np.array(mask)
52.             # 检查掩码的维度，确保是RGB（3通道）
53.             if len(mask.shape) != 3 or mask.shape[2] != 3:
54.                 raise ValueError(f"掩码维度错误: {mask.shape}, 期望为 (H,W,3)")
55.             mask_copy = np.zeros((mask.shape[0], mask.shape[1]), dtype=np.uint8)# 创建一个新的类别索引掩码
56.             # 将RGB颜色映射到类别索引
57.             # 遍历每种颜色，将对应像素设置为类别索引
58.             for k, color in enumerate(VOC_COLORMAP):
59.                 # 将每个颜色通道转换为布尔掩码
60.                 r_match = mask[:, :, 0] == color[0]
61.                 g_match = mask[:, :, 1] == color[1]
62.                 b_match = mask[:, :, 2] == color[2]
63.                 # 只有三个通道都匹配的像素才被分配为此类别
64.                 color_match = r_match & g_match & b_match
65.                 mask_copy[color_match] = k
66.             mask = torch.from_numpy(mask_copy).long() # 转换为PyTorch张量（长整型，用于交叉熵损失）
67.         return img, mask
68.     def __len__(self):
69.         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.2 FCN主干网络搭建

        在完成了数据加载函数之后我们便可以开始完成FCN主干网络的搭建了，我们将从最底子的FCN32s主干网络开始，逐步完成FCN16s和FCN8s，首先我们先创建fcn_model.py文件，同时导入一些必要的模块

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.2.1 完成双线性插值辅助函数

        在构建FCN网络时，我们必要用到反卷积（ConvTranspose2d）层来进行上采样，而这种上采样常常使用双线性插值初始化，以确保上采样时图像的平滑性。我们将在网络的_initialize_weights方法中使用我们自界说的双线性插值方法来初始化这些反卷积层的权重，。

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)

复制代码

        上面的代码确保了反卷积层的权重被正确初始化为适合图像上采样的双线性插值权重。接下来，我们实现_make_bilinear_weights方法用于生成双线性插值的权重矩阵。

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

复制代码

        通过这个方法，我们可以生成一个num_channels个通道的双线性插值权重矩阵，并将其赋值给反卷积层的权重。这样，网络在练习时，反卷积层将可以或许根据这些初始化的权重执行平滑的上采样操纵。
3.2.2 搭建FCN32s主干网络

        FCN32s是FCN系列中最底子的版本，其核心头脑是直接将VGG16的全连接层转换为卷积层，并通过32倍上采样规复原图分辨率。下面我将详细讲解如何用PyTorch实现这一布局，首先我们创建class FCN32s(nn.Module):网络

1. class FCN32s(nn.Module):
2.     def __init__(self, num_classes=NUM_CLASSES, pretrained=True):
3.         super(FCN32s, self).__init__()

复制代码

        FCN的特性提取主干网络使用的是VGG16，torchvision中有现成的网络架构可以直接导入，但是由于我们要实现分割，因此我们导入了VGG16后还必要将VGG16网络特性层做一些修改，使其可以或许适应全卷积网络布局。主要操纵便是使用Sequential操纵来提取VGG16中的卷积层。我们使用nn.Sequential将VGG16的前五个卷积块封装保留下来，接着我们将VGG16网络中的全连接层替换为1x1卷积层，最后我们通过一个上采样操纵将低分辨率的特性图规复到原始输入图像的尺寸，从而可以或许进行像素级别的分割预测。

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()

复制代码

        接着我们来实现这个网络的前向推理部分，在这部分里，我们首先记录输入图像的尺寸，input_size = x.size()[2:]，这是为了在经过上采样后确保输出的尺寸与输入图像一致。接下来，输入图像会依次通过VGG16网络中的卷积层进行特性提取。我们将图像通过features1到features5（分别对应VGG16中的五个卷积块）进行处理之后，我们会得到了一个较低分辨率的特性图，为了将这些低分辨率特性图转化为像素级的预测，我们接下来将它们通过两层1x1卷积（即fc6和fc7）进行处理，并使用ReLU激活函数进行非线性转换，同时为了防止过拟合我们在每一层后都应用了Dropout，接着经太过类层score后，我们便得到了一个终极的输出特性图，此中每个像素点的通道对应于一个类别的分割效果，然后我们便可以通过转置卷积（upsample）对输出进行32倍上采样，将特性图规复到原始图像的尺寸。最后，我们裁剪输出的尺寸，确保它与输入图像的大小一致。

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

复制代码

        完备的FCN32s网络如下：

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模型        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):
5.         # 初始化反卷积层的权重为双线性上采样
6.         for m in self.modules():
7.             if isinstance(m, nn.ConvTranspose2d):
8.                 # 双线性上采样的初始化
9.                 m.weight.data.zero_()
10.                 m.weight.data = self._make_bilinear_weights(m.kernel_size[0], m.out_channels)
11.     def _make_bilinear_weights(self, size, num_channels):
12.         """生成双线性插值的权重"""
13.         factor = (size + 1) // 2
14.         if size % 2 == 1:
15.             center = factor - 1
16.         else:
17.             center = factor - 0.5
18.         og = torch.FloatTensor(size, size)
19.         for i in range(size):
20.             for j in range(size):
21.                 og[i, j] = (1 - abs((i - center) / factor)) * (1 - abs((j - center) / factor))
22.         filter = torch.zeros(num_channels, num_channels, size, size)
23.         for i in range(num_channels):
24.             filter[i, i] = og
25.         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有变革的地方：

1. ### FCN16s
2. class FCN16s(nn.Module):
3.     def __init__(self, num_classes=NUM_CLASSES, pretrained=True):
4.         # 获取特征提取部分
5.         # 分段处理VGG16特征
7.         ######以上和FCN32s保持一致#########
8.         # 全连接层替换为1x1卷积
9.         self.fc6 = nn.Conv2d(512, 4096, kernel_size=7, padding=3)
10.         self.relu6 = nn.ReLU(inplace=True)
11.         self.drop6 = nn.Dropout2d()
12.         self.fc7 = nn.Conv2d(4096, 4096, kernel_size=1)
13.         self.relu7 = nn.ReLU(inplace=True)
14.         self.drop7 = nn.Dropout2d()
15.         self.score_fr = nn.Conv2d(4096, num_classes, kernel_size=1)# 分类层
16.         self.score_pool4 = nn.Conv2d(512, num_classes, kernel_size=1)# pool4的1x1卷积，用于特征融合
17.         # 2倍上采样conv7特征
18.         self.upsample2 = nn.ConvTranspose2d(num_classes, num_classes, kernel_size=4, stride=2, padding=1, bias=False)
19.         # 16倍上采样回原始图像大小
20.         self.upsample16 = nn.ConvTranspose2d(num_classes, num_classes, kernel_size=32, stride=16, padding=8, bias=False)
21.         # 初始化参数
22.         self._initialize_weights()
23.     def forward(self, x):
24.         input_size = x.size()[2:]# 记录输入尺寸用于上采样
25.         # 编码器 (VGG16)
26.         x = self.features1(x)
27.         x = self.features2(x)
28.         x = self.features3(x)
29.         pool4 = self.features4(x)# 保存pool4的输出用于后续融合
30.         x = self.features5(pool4)
31.         # 全连接层 (以卷积形式实现)
32.         x = self.relu6(self.fc6(x))
33.         x = self.drop6(x)
34.         x = self.relu7(self.fc7(x))
35.         x = self.drop7(x)
36.         x = self.score_fr(x)# 分类
37.         # 2倍上采样
38.         x = self.upsample2(x)
39.         # 获取pool4的分数并裁剪
40.         score_pool4 = self.score_pool4(pool4)
41.         score_pool4 = score_pool4[:, :, :x.size()[2], :x.size()[3]]
42.         x = x + score_pool4# 融合特征
43.         x = self.upsample16(x)# 16倍上采样回原始尺寸
44.         x = x[:, :, :input_size[0], :input_size[1]]# 裁剪到原始图像尺寸
45.         return x
47. ### FCN8s
48. class FCN8s(nn.Module):
49.     def __init__(self, num_classes=NUM_CLASSES, pretrained=True):
50.         # 获取特征提取部分
51.         # 分段处理VGG16特征
53.         ######以上和FCN32s保持一致#########
54.          # 全连接层替换为1x1卷积
55.         self.fc6 = nn.Conv2d(512, 4096, kernel_size=7, padding=3)
56.         self.relu6 = nn.ReLU(inplace=True)
57.         self.drop6 = nn.Dropout2d()
58.         self.fc7 = nn.Conv2d(4096, 4096, kernel_size=1)
59.         self.relu7 = nn.ReLU(inplace=True)
60.         self.drop7 = nn.Dropout2d()
61.         self.score_fr = nn.Conv2d(4096, num_classes, kernel_size=1)# 分类层
62.         # pool3和pool4的1x1卷积，用于特征融合
63.         self.score_pool4 = nn.Conv2d(512, num_classes, kernel_size=1)
64.         self.score_pool3 = nn.Conv2d(256, num_classes, kernel_size=1)
65.         # 2倍上采样conv7特征
66.         self.upsample2_1 = nn.ConvTranspose2d(num_classes, num_classes, kernel_size=4, stride=2, padding=1, bias=False)
67.         # 2倍上采样融合后的特征
68.         self.upsample2_2 = nn.ConvTranspose2d(num_classes, num_classes, kernel_size=4, stride=2, padding=1, bias=False)
69.         # 8倍上采样回原始图像大小
70.         self.upsample8 = nn.ConvTranspose2d(num_classes, num_classes, kernel_size=16, stride=8, padding=4, bias=False)
71.         # 初始化参数
72.         self._initialize_weights()
73.     def forward(self, x):
74.         input_size = x.size()[2:]# 记录输入尺寸用于上采样
75.         # 编码器 (VGG16)
76.         x = self.features1(x)
77.         x = self.features2(x)
78.         pool3 = self.features3(x)# 保存pool3的输出用于后续融合
79.         pool4 = self.features4(pool3)# 保存pool4的输出用于后续融合
80.         x = self.features5(pool4)
81.         # 全连接层 (以卷积形式实现)
82.         x = self.relu6(self.fc6(x))
83.         x = self.drop6(x)
84.         x = self.relu7(self.fc7(x))
85.         x = self.drop7(x)
86.         x = self.score_fr(x)# 分类
87.         x = self.upsample2_1(x)# 2倍上采样
88.         # 获取pool4的分数并裁剪
89.         score_pool4 = self.score_pool4(pool4)
90.         score_pool4 = score_pool4[:, :, :x.size()[2], :x.size()[3]]
91.         x = x + score_pool4 # 第一次融合特征 (pool5上采样 + pool4)
92.         x = self.upsample2_2(x)# 再次2倍上采样
93.         # 获取pool3的分数并裁剪
94.         score_pool3 = self.score_pool3(pool3)
95.         score_pool3 = score_pool3[:, :, :x.size()[2], :x.size()[3]]
96.         x = x + score_pool3# 第二次融合特征 (第一次融合的上采样 + pool3)
97.         x = self.upsample8(x)# 8倍上采样回原始尺寸
98.         x = x[:, :, :input_size[0], :input_size[1]]# 裁剪到原始图像尺寸
99.         return x

复制代码

3.3 完成练习脚本train.py

        在完成了FCN模型的网络布局搭建后，我们必要编写一个完备的练习脚本来对模型进行练习和评估。这个脚本将包含数据加载、模型练习、验证评估以及效果可视化等功能。下面我将详细讲解练习脚本的各个组成部分。
3.3.1 导入所需模块息争析命令行参数

        首先，我们必要导入必要的模块，并界说一个参数解析器，使得我们可以通过命令行机动地调整练习参数，在这里我们界说了一系列参数，包括数据集路径、模型范例选择、练习超参数（批大小、轮数、学习率等）以及查抄点相干参数，使得我们可以机动控制练习过程。

1. import os
2. import time
3. import argparse
4. import numpy as np
5. import torch
6. import torch.nn as nn
7. import torch.optim as optim
8. from torch.utils.data import DataLoader
9. from tqdm import tqdm
10. import matplotlib.pyplot as plt
11. import gc
12. from dataload import get_data_loaders, NUM_CLASSES, decode_segmap
13. from fcn_model import get_fcn_model
14. def parse_args():
15.     parser = argparse.ArgumentParser(description='FCN 语义分割 PyTorch 实现')
16.     parser.add_argument('--voc-root', type=str, default='',
17.                         help='VOC数据集根目录')
18.     parser.add_argument('--model-type', type=str, default='fcn8s', choices=['fcn8s', 'fcn16s', 'fcn32s'],
19.                         help='FCN模型类型 (fcn8s, fcn16s, fcn32s)')
20.     parser.add_argument('--batch-size', type=int, default=4,
21.                         help='训练的批次大小')
22.     parser.add_argument('--epochs', type=int, default=50,
23.                         help='训练的轮数')
24.     parser.add_argument('--lr', type=float, default=0.005,
25.                         help='学习率')
26.     parser.add_argument('--momentum', type=float, default=0.9,
27.                         help='SGD动量')
28.     parser.add_argument('--weight-decay', type=float, default=1e-4,
29.                         help='权重衰减')
30.     parser.add_argument('--num-workers', type=int, default=4,
31.                         help='数据加载线程数')
32.     parser.add_argument('--checkpoint-dir', type=str, default='checkpoints',
33.                         help='模型保存目录')
34.     parser.add_argument('--resume', type=str, default=None,
35.                         help='恢复训练的检查点路径')
36.     return parser.parse_args()

复制代码

3.3.2 模型评估函数

        接下来，我们界说一个评估函数，用于在验证集上评估模型性能。语义分割任务常用的评估指标包括像素正确率和平均交并比（mIoU），在这个评估函数中，我们通过遍历验证集的每个批次，盘算模型的预测效果与真实标签之间的损失、像素正确率以及每个类别的IoU（交并比）

1. def evaluate(model, val_loader, criterion, device):
2.     model.eval()
3.     total_loss = 0.0
4.     total_corrects = 0
5.     total_pixels = 0
6.     class_iou = np.zeros(NUM_CLASSES)
7.     class_pixels = np.zeros(NUM_CLASSES)
8.    
9.     with torch.no_grad():
10.         for images, targets in tqdm(val_loader, desc='Evaluation'):
11.             images = images.to(device)
12.             targets = targets.to(device)
13.             
14.             outputs = model(images)
15.             loss = criterion(outputs, targets)
16.             
17.             total_loss += loss.item() * images.size(0)
18.             
19.             _, preds = torch.max(outputs, 1)
20.             
21.             # 计算像素准确率
22.             correct = (preds == targets).sum().item()
23.             total_corrects += correct
24.             total_pixels += targets.numel()
25.             
26.             # 计算每个类别的IoU
27.             for cls in range(NUM_CLASSES):
28.                 pred_inds = preds == cls
29.                 target_inds = targets == cls
30.                 intersection = (pred_inds & target_inds).sum().item()
31.                 union = (pred_inds | target_inds).sum().item()
32.                
33.                 if union > 0:
34.                     class_iou[cls] += intersection / union
35.                     class_pixels[cls] += 1
36.             
37.             del images, targets, outputs, preds
38.             torch.cuda.empty_cache()  
39.    
40.     # 计算平均指标
41.     val_dataset_size = len(val_loader.dataset) if hasattr(val_loader.dataset, '__len__') else len(val_loader) * val_loader.batch_size
42.     avg_loss = total_loss / val_dataset_size
43.     pixel_acc = total_corrects / total_pixels
44.    
45.     # 计算每个类别的平均IoU
46.     for cls in range(NUM_CLASSES):
47.         if class_pixels[cls] > 0:
48.             class_iou[cls] /= class_pixels[cls]
49.    
50.     # 计算mIoU (平均交并比)
51.     miou = np.mean(class_iou)
52.    
53.     gc.collect()
54.     torch.cuda.empty_cache()
55.    
56.     return avg_loss, pixel_acc, miou, class_iou

复制代码

3.3.3 分割效果可视化函数

        为了直观地观察分割效果，我们界说一个函数用于保存预测效果的可视化图像。这个函数将从验证集中取出一定命量的样本，通过模型进行预测，然后将原始图像、真实标签和预测效果并排显示并保存为图像文件，这样我们就可以直观地观察模型的分割效果

1. def save_predictions(model, val_loader, device, output_dir='outputs', num_samples=5):
2.     if not os.path.exists(output_dir):
3.         os.makedirs(output_dir)
4.    
5.     model.eval()
6.     with torch.no_grad():
7.         for i, (images, targets) in enumerate(val_loader):
8.             if i >= num_samples:
9.                 break
10.             images = images.to(device)
11.             targets = targets.to(device)
12.             outputs = model(images)
13.             _, preds = torch.max(outputs, 1)
14.             # 转换为NumPy数组用于可视化
15.             images_np = images.cpu().numpy()
16.             targets_np = targets.cpu().numpy()
17.             preds_np = preds.cpu().numpy()
18.             # 对每个样本进行可视化
19.             for b in range(images.size(0)):
20.                 if b >= 3:  # 限制每个批次只保存前3个样本
21.                     break
22.                 img = images_np[b].transpose(1, 2, 0)
23.                 mean = np.array([0.485, 0.456, 0.406])
24.                 std = np.array([0.229, 0.224, 0.225])
25.                 img = img * std + mean
26.                 img = np.clip(img, 0, 1)
27.                 target_rgb = decode_segmap(targets_np[b])
28.                 pred_rgb = decode_segmap(preds_np[b])
29.                 plt.figure(figsize=(15, 5))
30.                 plt.subplot(1, 3, 1)
31.                 plt.title('Input Image')
32.                 plt.imshow(img)
33.                 plt.axis('off')
34.                 plt.subplot(1, 3, 2)
35.                 plt.title('Ground Truth')
36.                 plt.imshow(target_rgb)
37.                 plt.axis('off')
38.                 plt.subplot(1, 3, 3)
39.                 plt.title('Prediction')
40.                 plt.imshow(pred_rgb)
41.                 plt.axis('off')
42.                 plt.tight_layout()
43.                 plt.savefig(os.path.join(output_dir, f'sample_{i}_{b}.png'))
44.                 plt.close()
46.             del images, targets, outputs, preds
47.             torch.cuda.empty_cache()
48.    
49.     gc.collect()
50.     torch.cuda.empty_cache()

复制代码

3.3.4 主练习函数

        最后，我们编写主函数，实现完备的练习流程。主函数是整个练习脚本的核心，它将各个组件有机地整合在一起，形成完备的练习流程。首先，它通过parse_args()解析命令行输入的各项参数，如数据集路径、模型范例、批量大小等，使练习过程更加机动可控。之后，它会调用get_data_loaders()函数加载并预处理VOC数据集，同时创建数据加载器以便批量获取练习和验证样本。接着，根据参入参数指定的模型范例（FCN8s/FCN16s/FCN32s）实例化相应的网络布局，并将其转移到可用的盘算设备（GPU或CPU）上。在优化策略方面，主函数使用交叉熵损失函数（忽略255标签值）评估分割质量，采用动员量的SGD优化器更新网络参数，并通过学习率调度器在练习后期低落学习率以得到更精致的优化效果。假如参入参数提供了查抄点路径，函数会从中规复模型权重、优化器状态和练习进度，实现断点续训。核心的练习循环涵盖了完备的练习-评估-保存流程：每个epoch内先在练习集上进行前向传播、损失盘算、反向传播和参数更新；然后在验证集上评估模型性能（损失值、像素正确率和mIoU）；当取得更高mIoU时，保存最佳模型并生成可视化效果，同时定期保存最新模型以防练习停止。练习完成后，主函数会绘制整个练习过程的损失曲线、正确率曲线和mIoU曲线，直观展示模型的学习轨迹和性能变革，资助大家更好地理解练习动态并优化练习策略。

1. def main():
2.     args = parse_args()
3.     if not os.path.exists(args.checkpoint_dir):
4.         os.makedirs(args.checkpoint_dir)
5.     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
6.     print(f'使用设备: {device}')
7.     # 加载数据
8.     train_loader, val_loader = get_data_loaders(
9.         args.voc_root,
10.         batch_size=args.batch_size,
11.         num_workers=args.num_workers
12.     )
13.     train_dataset_size = len(train_loader.dataset) if hasattr(train_loader.dataset, '__len__') else len(train_loader) * train_loader.batch_size
14.     val_dataset_size = len(val_loader.dataset) if hasattr(val_loader.dataset, '__len__') else len(val_loader) * val_loader.batch_size
15.     print(f'训练样本数: {train_dataset_size}, 验证样本数: {val_dataset_size}')
16.     # 创建模型
17.     model = get_fcn_model(model_type=args.model_type, num_classes=NUM_CLASSES, pretrained=True)
18.     model = model.to(device)
19.     # 定义损失函数和优化器
20.     criterion = nn.CrossEntropyLoss(ignore_index=255)  
21.     optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)
22.     scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)
23.     # 恢复训练
24.     start_epoch = 0
25.     best_miou = 0.0
26.     if args.resume:
27.         if os.path.isfile(args.resume):
28.             print(f'加载检查点: {args.resume}')
29.             checkpoint = torch.load(args.resume)
30.             start_epoch = checkpoint['epoch']
31.             best_miou = checkpoint['best_miou']
32.             model.load_state_dict(checkpoint['model_state_dict'])
33.             optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
34.             scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
35.             print(f'从 epoch {start_epoch} 恢复训练, 最佳 mIoU: {best_miou:.4f}')
36.         else:
37.             print(f'找不到检查点: {args.resume}')
38.     # 训练循环
39.     print(f'开始训练 {args.model_type} 模型, 共 {args.epochs} 轮...')
40.     # 记录训练历史
41.     history = {
42.         'train_loss': [],
43.         'val_loss': [],
44.         'pixel_acc': [],
45.         'miou': []
46.     }
47.     for epoch in range(start_epoch, args.epochs):
48.         # 训练阶段
49.         model.train()
50.         train_loss = 0.0
51.         batch_count = 0
52.         
53.         t0 = time.time()
54.         for images, targets in tqdm(train_loader, desc=f'Epoch {epoch+1}/{args.epochs}'):
55.             images = images.to(device)
56.             targets = targets.to(device)
57.             
58.             optimizer.zero_grad()
59.             
60.             outputs = model(images)
61.             loss = criterion(outputs, targets)
62.             
63.             loss.backward()
64.             optimizer.step()
65.             
66.             train_loss += loss.item() * images.size(0)
67.             batch_count += 1
68.             
69.             del images, targets, outputs, loss
70.             
71.             if batch_count % 10 == 0:
72.                 torch.cuda.empty_cache()
73.         
74.         train_loss = train_loss / train_dataset_size
75.         history['train_loss'].append(train_loss)
76.         # 调整学习率
77.         scheduler.step()
78.         gc.collect()
79.         torch.cuda.empty_cache()
80.         # 评估模型
81.         val_loss, pixel_acc, miou, class_iou = evaluate(model, val_loader, criterion, device)
82.         history['val_loss'].append(val_loss)
83.         history['pixel_acc'].append(pixel_acc)
84.         history['miou'].append(miou)
85.         
86.         # 打印进度
87.         epoch_time = time.time() - t0
88.         print(f'Epoch {epoch+1}/{args.epochs} - '
89.               f'Time: {epoch_time:.2f}s - '
90.               f'Train Loss: {train_loss:.4f} - '
91.               f'Val Loss: {val_loss:.4f} - '
92.               f'Pixel Acc: {pixel_acc:.4f} - '
93.               f'mIoU: {miou:.4f}')
94.         
95.         # 保存最佳模型
96.         if miou > best_miou:
97.             best_miou = miou
98.             torch.save({
99.                 'epoch': epoch + 1,
100.                 'model_state_dict': model.state_dict(),
101.                 'optimizer_state_dict': optimizer.state_dict(),
102.                 'scheduler_state_dict': scheduler.state_dict(),
103.                 'best_miou': best_miou,
104.             }, os.path.join(args.checkpoint_dir, f'{args.model_type}_best.pth'))
105.             print(f'保存最佳模型, mIoU: {best_miou:.4f}')
106.             
107.             # 生成可视化结果
108.             save_predictions(model, val_loader, device)
109.         
110.         # 保存最新模型
111.         torch.save({
112.             'epoch': epoch + 1,
113.             'model_state_dict': model.state_dict(),
114.             'optimizer_state_dict': optimizer.state_dict(),
115.             'scheduler_state_dict': scheduler.state_dict(),
116.             'best_miou': best_miou,
117.         }, os.path.join(args.checkpoint_dir, f'{args.model_type}_latest.pth'))
118.         
119.         gc.collect()
120.         torch.cuda.empty_cache()
121.     plt.figure(figsize=(12, 10))
122.     plt.subplot(2, 2, 1)
123.     plt.plot(history['train_loss'], label='Train')
124.     plt.plot(history['val_loss'], label='Validation')
125.     plt.title('Loss')
126.     plt.xlabel('Epoch')
127.     plt.ylabel('Loss')
128.     plt.legend()
129.     plt.subplot(2, 2, 2)
130.     plt.plot(history['pixel_acc'])
131.     plt.title('Pixel Accuracy')
132.     plt.xlabel('Epoch')
133.     plt.ylabel('Accuracy')
134.     plt.subplot(2, 2, 3)
135.     plt.plot(history['miou'])
136.     plt.title('Mean IoU')
137.     plt.xlabel('Epoch')
138.     plt.ylabel('mIoU')
139.     plt.tight_layout()
140.     plt.savefig(os.path.join(args.checkpoint_dir, f'{args.model_type}_history.png'))
141.     plt.close()
142.     print(f'训练完成! 最佳 mIoU: {best_miou:.4f}')

复制代码

        完成了练习代码之后我们便可以开始练习啦！我们输入以下命令即可开始练习：

1. 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参数则可以选择是否将掩码叠加在原图上面

1. import os
2. import argparse
3. import numpy as np
4. import torch
5. import torchvision.transforms as transforms
6. from PIL import Image
7. import matplotlib.pyplot as plt
9. from dataload import VOC_CLASSES, VOC_COLORMAP, NUM_CLASSES, decode_segmap
10. from fcn_model import get_fcn_model
12. def parse_args():
13.     parser = argparse.ArgumentParser(description='FCN语义分割模型预测')
14.     parser.add_argument('--model-path', type=str, required=True,
15.                         help='预训练模型路径')
16.     parser.add_argument('--model-type', type=str, default='fcn8s', choices=['fcn8s', 'fcn16s', 'fcn32s'],
17.                         help='FCN模型类型 (fcn8s, fcn16s, fcn32s)')
18.     parser.add_argument('--image-path', type=str, required=True,
19.                         help='输入图像路径，可以是单个图像或者目录')
20.     parser.add_argument('--output-dir', type=str, default='results',
21.                         help='结果保存目录')
22.     parser.add_argument('--overlay', action='store_true',
23.                         help='是否将分割结果与原图叠加')
24.     parser.add_argument('--no-cuda', action='store_true',
25.                         help='禁用CUDA')
26.     return parser.parse_args()

复制代码

3.4.2 图像预处理和后处理函数

        接下来，我们界说两个辅助函数：一个用于预处理输入图像，使其符合模型的输入要求；另一个用于将分割效果与原图叠加，加强可视化效果。

1. def preprocess_image(image_path):
2.     """预处理输入图像"""
3.     image = Image.open(image_path).convert('RGB')
4.     # 图像预处理变换
5.     transform = transforms.Compose([
6.         transforms.Resize(320),
7.         transforms.CenterCrop(320),
8.         transforms.ToTensor(),
9.         transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
10.     ])
11.     input_tensor = transform(image)
12.     input_batch = input_tensor.unsqueeze(0)  
14.     return input_batch, image
16. def overlay_segmentation(image, segmentation, alpha=0.7):
17.     """将分割结果与原图叠加"""
18.     image_np = np.array(image)
19.     segmentation_resized = np.array(Image.fromarray(segmentation.astype(np.uint8)).resize(
20.         (image_np.shape[1], image_np.shape[0]), Image.NEAREST))
21.     overlay = image_np.copy()
22.     for i in range(3):
23.         overlay[:, :, i] = image_np[:, :, i] * (1 - alpha) + segmentation_resized[:, :, i] * alpha
24.     return overlay.astype(np.uint8)

复制代码

        preprocess_image函数将输入图像调整为同一大小（320×320），转换为张量格式，并应用ImageNet数据集的尺度归一化。overlay_segmentation函数则担当原始图像和分割图，按指定的透明度（默认0.7）将它们叠加在一起，使得分割效果更直观。
3.4.3 预测及可视化函数

        下面我们实现对图像进行预测和可视化的功能：

1. def predict_image(model, image_path, device, overlay=False):
2.     """对单个图像进行预测"""
3.     input_batch, original_image = preprocess_image(image_path)
4.     input_batch = input_batch.to(device)
5.     with torch.no_grad():
6.         output = model(input_batch)
7.         output = torch.nn.functional.softmax(output, dim=1)
8.         _, pred = torch.max(output, 1)
9.         pred = pred.cpu().numpy()[0]  
10.     # 将预测结果转换为彩色分割图
11.     segmentation_map = decode_segmap(pred)
12.     if overlay:
13.         result = overlay_segmentation(original_image, segmentation_map)
14.     else:
15.         result = segmentation_map
16.    
17.     return result, pred, original_image
19. def predict_and_visualize(model, image_path, output_dir, device, overlay=False):
20.     """预测图像并可视化结果"""
21.     # 如果图像路径是目录
22.     if os.path.isdir(image_path):
23.         image_files = [f for f in os.listdir(image_path) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
24.         for image_file in image_files:
25.             file_path = os.path.join(image_path, image_file)
26.             visualize_prediction(model, file_path, output_dir, device, overlay)
27.     else:
28.         visualize_prediction(model, image_path, output_dir, device, overlay)
30. def visualize_prediction(model, image_path, output_dir, device, overlay=False):
31.     """可视化单个图像的预测结果"""
32.     os.makedirs(output_dir, exist_ok=True)
33.     # 预测图像
34.     result, pred, original_image = predict_image(model, image_path, device, overlay)
35.     # 保存结果
36.     base_name = os.path.basename(image_path).split('.')[0]
37.     plt.figure(figsize=(15, 5))
38.     plt.subplot(1, 3, 1)
39.     plt.title('原始图像')
40.     plt.imshow(original_image)
41.     plt.axis('off')
42.     plt.subplot(1, 3, 2)
43.     plt.title('分割结果')
44.     plt.imshow(decode_segmap(pred))
45.     plt.axis('off')
46.     plt.subplot(1, 3, 3)
47.     if overlay:
48.         plt.title('叠加结果')
49.     else:
50.         plt.title('分割结果')
51.     plt.imshow(result)
52.     plt.axis('off')
53.     plt.tight_layout()
54.     plt.savefig(os.path.join(output_dir, f'{base_name}_result.png'))
55.     plt.close()
56.     class_pixels = {}
57.     for i, class_name in enumerate(VOC_CLASSES):
58.         num_pixels = np.sum(pred == i)
59.         if num_pixels > 0:
60.             class_pixels[class_name] = num_pixels
61.     # 创建类别分布饼图
62.     if class_pixels:
63.         plt.figure(figsize=(10, 10))
64.         labels = list(class_pixels.keys())
65.         sizes = list(class_pixels.values())
66.         plt.pie(sizes, labels=labels, autopct='%1.1f%%', shadow=True, startangle=90)
67.         plt.axis('equal')
68.         plt.title('类别分布')
69.         plt.savefig(os.path.join(output_dir, f'{base_name}_class_dist.png'))
70.         plt.close()
71.     # 保存单独的分割图
72.     segmentation_img = Image.fromarray(decode_segmap(pred))
73.     segmentation_img.save(os.path.join(output_dir, f'{base_name}_segmentation.png'))

复制代码

        预测函数将首先加载并预处理图像，然后通过模型进行前向传播，获取预测效果。而预测效果先通过softmax转换为概率分布，然后选取概率最高的类别作为终极预测。最后，根据是否必要叠加展示，返回相应的可视化效果。
3.4.4 主函数实现

        最后，我们实现主函数，将全部功能整合起来，主函数首先解析命令行参数，然后根据参数创建相应的FCN模型。在加载预练习权重时，我特别思量了PyTorch差别版本的兼容性问题，使用了try-except布局来适应差别版本的加载方式。加载完模型后，将其设置为评估模式，然后调用预测和可视化函数处理指定的图像或图像目录。

1. def main():
2.     args = parse_args()
3.     device = torch.device('cuda' if torch.cuda.is_available() and not args.no_cuda else 'cpu')
4.     print(f'使用设备: {device}')
5.     # 创建模型
6.     model = get_fcn_model(model_type=args.model_type, num_classes=NUM_CLASSES, pretrained=False)
7.     checkpoint = torch.load(args.model_path, map_location=device)
8.    
9.     if 'model_state_dict' in checkpoint:
10.         model.load_state_dict(checkpoint['model_state_dict'])
11.         print(f'加载检查点: Epoch {checkpoint["epoch"]}, mIoU {checkpoint["best_miou"]:.4f}')
12.     else:
13.         model.load_state_dict(checkpoint)
14.         print(f'加载模型权重成功')
15.    
16.     model = model.to(device)
17.     model.eval()
18.    
19.     predict_and_visualize(model, args.image_path, args.output_dir, device, args.overlay)
20.     print(f'结果已保存到: {args.output_dir}')

复制代码

        完成了预测推理代码之后我们便可以使用如下命令进行推理：

1. python predict.py --model-path checkpoints/fcn8s_best.pth --image-path test_images/test.jpg --overlay

复制代码

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

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

免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！更多信息从访问主页：qidao123.com:ToB企服之家，中国第一个企服评测及商务社交产业平台。