别再死记FCN结构了！用PyTorch从VGG16开始，手把手复现FCN-8s（附代码）

从VGG16到FCN-8s用PyTorch实现语义分割的终极实践指南在计算机视觉领域语义分割一直是极具挑战性的任务之一。不同于简单的图像分类语义分割需要模型对图像中的每一个像素进行分类这要求模型不仅要理解图像的整体内容还要捕捉到局部的细节特征。全卷积网络FCN作为这一领域的里程碑式工作首次展示了如何将传统的分类网络改造为能够处理任意尺寸输入的像素级预测模型。本文将带您从零开始基于PyTorch框架从预训练的VGG16出发一步步构建完整的FCN-8s模型。1. 理解FCN的核心思想FCN的成功在于它巧妙地解决了传统CNN在分割任务中的几个关键限制。传统CNN通常以固定尺寸的输入图像开始经过一系列卷积和池化操作后最终通过全连接层输出分类结果。这种结构存在两个主要问题固定尺寸输入全连接层的存在要求输入尺寸必须固定因为全连接层的权重矩阵大小是预先确定的。空间信息丢失全连接层将特征图展平为一维向量丢失了原有的空间结构信息。FCN通过以下创新解决了这些问题全卷积化将最后的全连接层替换为1×1卷积层使网络可以接受任意尺寸的输入上采样通过转置卷积反卷积操作将低分辨率预测图上采样到输入图像尺寸跳跃连接融合来自不同深度的特征图结合高层语义信息和低层细节信息# 传统CNN与FCN结构对比示意代码 class TraditionalCNN(nn.Module): def __init__(self): super().__init__() self.conv_layers nn.Sequential(...) # 卷积层 self.fc_layers nn.Sequential(...) # 全连接层 class FCN(nn.Module): def __init__(self): super().__init__() self.conv_layers nn.Sequential(...) # 卷积层 self.conv1x1 nn.Conv2d(...) # 1x1卷积代替全连接 self.upsample nn.ConvTranspose2d(...) # 转置卷积上采样2. 准备VGG16骨干网络VGG16作为FCN论文中表现最好的骨干网络其规整的结构和较深的层次非常适合作为特征提取器。我们将使用PyTorch中预训练的VGG16模型作为起点。关键步骤加载预训练VGG16模型保留特征提取部分卷积层移除最后的全连接层和分类器识别并标记需要用于跳跃连接的特征层import torchvision.models as models def get_vgg16_features(pretrainedTrue): vgg16 models.vgg16(pretrainedpretrained).features # 标记需要保留用于跳跃连接的层 layer_indices { pool3: 17, # 用于FCN-8s的1/8特征 pool4: 24, # 用于FCN-16s的1/16特征 pool5: 31 # 最后的1/32特征 } return vgg16, layer_indicesVGG16特征提取层结构层类型输出尺寸 (输入224×224)说明ConvReLU224×224×64第一层卷积MaxPool112×112×64第一次下采样ConvReLU112×112×128第二层卷积MaxPool56×56×128第二次下采样ConvReLU56×56×256第三层卷积MaxPool28×28×256第三次下采样pool3ConvReLU28×28×512第四层卷积MaxPool14×14×512第四次下采样pool4ConvReLU14×14×512第五层卷积MaxPool7×7×512第五次下采样pool53. 构建FCN-8s网络结构FCN-8s之所以得名是因为它最终将特征图上采样8倍得到与输入相同尺寸的分割结果。相比FCN-32s和FCN-16sFCN-8s通过融合更多低层特征能够保留更多细节信息。3.1 全卷积化改造首先需要将VGG16的全连接层替换为1×1卷积层。在PyTorch中VGG16的最后一个全连接层输入维度是250887×7×512输出维度是4096。class FCN8s(nn.Module): def __init__(self, num_classes, pretrainedTrue): super().__init__() # 获取VGG16特征提取器和关键层索引 self.vgg16, self.layer_indices get_vgg16_features(pretrained) # 用1x1卷积代替全连接层 self.conv6 nn.Conv2d(512, 4096, kernel_size7, padding3) self.conv7 nn.Conv2d(4096, 4096, kernel_size1) self.score_fr nn.Conv2d(4096, num_classes, kernel_size1) # 用于跳跃连接的分支 self.score_pool3 nn.Conv2d(256, num_classes, kernel_size1) self.score_pool4 nn.Conv2d(512, num_classes, kernel_size1) # 上采样层 self.upscore2 nn.ConvTranspose2d( num_classes, num_classes, kernel_size4, stride2, padding1) self.upscore8 nn.ConvTranspose2d( num_classes, num_classes, kernel_size16, stride8, padding4) self.upscore_pool4 nn.ConvTranspose2d( num_classes, num_classes, kernel_size4, stride2, padding1)3.2 跳跃连接实现跳跃连接是FCN-8s性能优于FCN-32s和FCN-16s的关键。它通过将深层的高语义信息与浅层的细节信息相结合实现了更精确的分割边界。def forward(self, x): # 原始图像尺寸 h, w x.shape[2], x.shape[3] # 前向传播获取各层特征 pool3 None pool4 None for idx, layer in enumerate(self.vgg16): x layer(x) if idx self.layer_indices[pool3]: pool3 x # 1/8特征 elif idx self.layer_indices[pool4]: pool4 x # 1/16特征 # 主分支处理 x F.relu(self.conv6(x)) x F.dropout(x, 0.5) x F.relu(self.conv7(x)) x F.dropout(x, 0.5) x self.score_fr(x) # 第一次上采样(32s) x self.upscore2(x) # 1/16 # 融合pool4特征(16s) pool4_score self.score_pool4(pool4) x x[:, :, :pool4_score.size(2), :pool4_score.size(3)] pool4_score # 第二次上采样(16s) x self.upscore_pool4(x) # 1/8 # 融合pool3特征(8s) pool3_score self.score_pool3(pool3) x x[:, :, :pool3_score.size(2), :pool3_score.size(3)] pool3_score # 最终上采样到原始尺寸 x self.upscore8(x) x x[:, :, :h, :w] # 裁剪到原始尺寸 return x注意在实际实现中跳跃连接的特征图尺寸可能不完全匹配需要进行适当的裁剪或填充。这里我们采用了裁剪方式确保相加操作能够顺利进行。4. 训练策略与技巧训练FCN网络需要特别注意几个关键点包括损失函数选择、学习率设置和数据增强策略。4.1 损失函数与评估指标语义分割任务通常使用像素级的交叉熵损失函数。由于分割任务中类别分布往往不均衡例如背景像素远多于前景像素可以考虑使用加权交叉熵或Dice损失。def train_step(model, optimizer, data, target, criterion): optimizer.zero_grad() output model(data) loss criterion(output, target) loss.backward() optimizer.step() return loss.item() # 使用加权交叉熵损失 class_weight torch.tensor([1.0, 2.0, 3.0]) # 示例权重 criterion nn.CrossEntropyLoss(weightclass_weight)常用评估指标像素准确率Pixel Accuracy正确分类的像素比例平均交并比mIoU各类别IoU的平均值频率加权IoU考虑类别频率的加权IoU4.2 数据增强与预处理由于FCN可以接受任意尺寸的输入我们可以使用更灵活的数据增强策略随机缩放0.5-2.0倍随机水平/垂直翻转颜色抖动亮度、对比度、饱和度调整随机旋转-10°到10°随机裁剪保持长宽比from torchvision import transforms train_transform transforms.Compose([ transforms.RandomResizedCrop(224, scale(0.5, 2.0)), transforms.RandomHorizontalFlip(), transforms.ColorJitter(0.2, 0.2, 0.2), transforms.RandomRotation(10), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ])4.3 训练参数设置基于原始FCN论文的建议我们可以采用以下训练策略优化器带动量的SGD初始学习率1e-3动量0.9权重衰减5e-4批次大小根据GPU内存选择通常8-16学习率调度在验证指标平稳时降低学习率optimizer torch.optim.SGD(model.parameters(), lr1e-3, momentum0.9, weight_decay5e-4) scheduler torch.optim.lr_scheduler.ReduceLROnPlateau( optimizer, max, patience3, factor0.1)5. 模型测试与可视化训练完成后我们需要对模型进行测试并可视化分割结果以评估其实际表现。5.1 测试流程测试时需要关闭dropout和batch normalization的随机性并确保输入图像经过与训练时相同的预处理。def evaluate(model, dataloader, device): model.eval() total_correct 0 total_pixels 0 with torch.no_grad(): for data, target in dataloader: data, target data.to(device), target.to(device) output model(data) pred output.argmax(dim1) total_correct (pred target).sum().item() total_pixels target.numel() accuracy total_correct / total_pixels return accuracy5.2 结果可视化将模型预测的分割结果与真实标注进行对比可视化可以直观地评估模型性能。import matplotlib.pyplot as plt def visualize_results(image, gt_mask, pred_mask): plt.figure(figsize(15, 5)) plt.subplot(1, 3, 1) plt.imshow(image.permute(1, 2, 0).cpu().numpy()) plt.title(Input Image) plt.subplot(1, 3, 2) plt.imshow(gt_mask.cpu().numpy(), cmapjet) plt.title(Ground Truth) plt.subplot(1, 3, 3) plt.imshow(pred_mask.argmax(dim0).cpu().numpy(), cmapjet) plt.title(Prediction) plt.show()5.3 不同版本FCN对比为了展示FCN-8s的优势我们可以将其与FCN-32s和FCN-16s进行对比模型版本上采样倍数跳跃连接优点缺点FCN-32s32×无实现简单细节丢失严重FCN-16s16×pool4保留部分细节中等精度FCN-8s8×pool3pool4细节保留最好实现较复杂在实际项目中我发现FCN-8s的分割边缘明显更加精细特别是在处理复杂场景时如街景中的行人、车辆等小目标。然而这种改进是以增加计算量为代价的需要根据具体应用场景进行权衡。