从V1到V3+：手把手带你复现DeepLab系列的核心模块（PyTorch代码详解）

从V1到V3+：手把手带你复现DeepLab系列的核心模块（PyTorch代码详解）

语义分割作为计算机视觉领域的核心任务之一，其目标是为图像中的每个像素分配语义标签。DeepLab系列模型凭借其创新的设计理念和卓越的性能表现，成为该领域的标杆性工作。本文将聚焦代码实践，通过PyTorch实现DeepLab各版本的核心模块，帮助开发者深入理解其技术演进脉络。

1. 环境准备与基础架构

在开始复现之前，我们需要搭建基础开发环境。推荐使用Python 3.8+和PyTorch 1.10+版本，这些版本能够很好地支持后续的空洞卷积等特性。

import torch import torch.nn as nn import torch.nn.functional as F from typing import List, Optional print(f"PyTorch版本: {torch.__version__}") print(f"CUDA可用: {torch.cuda.is_available()}")

DeepLab系列的基础架构通常基于修改后的ResNet或VGG网络。以下是一个基础的特征提取模块实现：

class BasicBlock(nn.Module): def __init__(self, in_channels, out_channels, stride=1, dilation=1): super().__init__() self.conv1 = nn.Conv2d( in_channels, out_channels, kernel_size=3, stride=stride, padding=dilation, dilation=dilation, bias=False ) self.bn1 = nn.BatchNorm2d(out_channels) self.conv2 = nn.Conv2d( out_channels, out_channels, kernel_size=3, padding=dilation, dilation=dilation, bias=False ) self.bn2 = nn.BatchNorm2d(out_channels) if stride != 1 or in_channels != out_channels: self.shortcut = nn.Sequential( nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride, bias=False), nn.BatchNorm2d(out_channels) ) else: self.shortcut = nn.Identity() def forward(self, x): out = F.relu(self.bn1(self.conv1(x))) out = self.bn2(self.conv2(out)) out += self.shortcut(x) return F.relu(out)

注意：在实际实现中，output_stride（输出步长）是一个关键参数，它决定了网络最终特征图相对于输入图像的下采样率。通常设置为16或8，需要在网络设计时统一考虑。

2. DeepLabV1核心：空洞卷积实现

DeepLabV1首次将空洞卷积引入语义分割任务，解决了传统CNN下采样导致的信息丢失问题。以下是空洞卷积的PyTorch实现：

class AtrousConv(nn.Module): def __init__(self, in_channels, out_channels, dilation): super().__init__() self.conv = nn.Conv2d( in_channels, out_channels, kernel_size=3, padding=dilation, dilation=dilation, bias=False ) self.bn = nn.BatchNorm2d(out_channels) def forward(self, x): return F.relu(self.bn(self.conv(x)))

为了验证空洞卷积的效果，我们可以对比普通卷积和空洞卷积的感受野：

DeepLabV1的网络结构调整策略包括：

• 将最后两个max-pool层的步长改为1，避免过度下采样
• 在高层网络中使用空洞卷积扩大感受野
• 最终输出通过双线性插值上采样8倍得到分割结果

3. DeepLabV2突破：ASPP模块详解

DeepLabV2提出了ASPP（Atrous Spatial Pyramid Pooling）模块，通过并行使用不同空洞率的卷积来捕获多尺度信息。以下是完整的ASPP实现：

class ASPP(nn.Module): def __init__(self, in_channels, out_channels=256, rates=[6, 12, 18]): super().__init__() modules = [] # 1×1卷积分支 modules.append(nn.Sequential( nn.Conv2d(in_channels, out_channels, 1, bias=False), nn.BatchNorm2d(out_channels), nn.ReLU() )) # 多尺度空洞卷积分支 for rate in rates: modules.append(nn.Sequential( nn.Conv2d(in_channels, out_channels, 3, padding=rate, dilation=rate, bias=False), nn.BatchNorm2d(out_channels), nn.ReLU() )) # 全局平均池化分支 modules.append(nn.Sequential( nn.AdaptiveAvgPool2d(1), nn.Conv2d(in_channels, out_channels, 1, bias=False), nn.BatchNorm2d(out_channels), nn.ReLU(), nn.Upsample(scale_factor=16, mode='bilinear', align_corners=True) )) self.branches = nn.ModuleList(modules) self.project = nn.Sequential( nn.Conv2d(out_channels * (len(rates)+2), out_channels, 1, bias=False), nn.BatchNorm2d(out_channels), nn.ReLU(), nn.Dropout(0.5) ) def forward(self, x): size = x.shape[-2:] features = [] for branch in self.branches: if isinstance(branch[-1], nn.Upsample): # 处理全局池化分支 feat = branch(x) else: feat = branch(x) features.append(feat) # 调整全局池化分支的大小 features[-1] = F.interpolate(features[-1], size=size, mode='bilinear', align_corners=True) x = torch.cat(features, dim=1) return self.project(x)

ASPP模块中各分支的作用：

• 1×1卷积：捕获原始尺度特征
• 多尺度空洞卷积：捕获不同感受野下的上下文信息
• 全局平均池化：提供图像级全局上下文

提示：在实际应用中，空洞率的选择需要根据output_stride进行调整。当output_stride=16时，常用rates=[6,12,18]；当output_stride=8时，rates应相应减半。

4. DeepLabV3改进：Multi-Grid策略与增强型ASPP

DeepLabV3引入了Multi-Grid策略来进一步优化空洞卷积的使用。以下是带有Multi-Grid的残差块实现：

class Bottleneck(nn.Module): expansion = 4 def __init__(self, in_channels, out_channels, stride=1, dilation=1, multi_grid=(1,1,1)): super().__init__() width = out_channels // self.expansion self.conv1 = nn.Conv2d(in_channels, width, 1, bias=False) self.bn1 = nn.BatchNorm2d(width) # 使用multi_grid调整各层的空洞率 self.conv2 = nn.ModuleList() for mg in multi_grid: self.conv2.append(nn.Sequential( nn.Conv2d(width, width, 3, stride=stride, padding=dilation*mg, dilation=dilation*mg, bias=False), nn.BatchNorm2d(width), nn.ReLU(inplace=True) )) self.conv3 = nn.Conv2d(width, out_channels, 1, bias=False) self.bn3 = nn.BatchNorm2d(out_channels) if stride != 1 or in_channels != out_channels: self.shortcut = nn.Sequential( nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride, bias=False), nn.BatchNorm2d(out_channels) ) else: self.shortcut = nn.Identity() def forward(self, x): identity = self.shortcut(x) out = F.relu(self.bn1(self.conv1(x))) for conv in self.conv2: out = conv(out) out = self.bn3(self.conv3(out)) out += identity return F.relu(out)

DeepLabV3对ASPP的主要改进包括：

1. 在ASPP中增加了Batch Normalization
2. 引入了图像级特征（全局平均池化）
3. 移除了CRF后处理

以下是改进后的ASPP模块参数配置建议：

5. DeepLabV3+创新：编码器-解码器结构与深度可分离卷积

DeepLabV3+最大的改进是引入了编码器-解码器结构和深度可分离卷积。以下是解码器模块的实现：

class Decoder(nn.Module): def __init__(self, low_level_channels, num_classes): super().__init__() self.conv1 = nn.Conv2d(low_level_channels, 48, 1, bias=False) self.bn1 = nn.BatchNorm2d(48) self.last_conv = nn.Sequential( nn.Conv2d(304, 256, 3, padding=1, bias=False), nn.BatchNorm2d(256), nn.ReLU(), nn.Dropout(0.5), nn.Conv2d(256, 256, 3, padding=1, bias=False), nn.BatchNorm2d(256), nn.ReLU(), nn.Dropout(0.1), nn.Conv2d(256, num_classes, 1) ) def forward(self, x, low_level_feat): low_level_feat = self.conv1(low_level_feat) low_level_feat = self.bn1(low_level_feat) low_level_feat = F.relu(low_level_feat) # 调整低层特征图尺寸 x = F.interpolate(x, size=low_level_feat.shape[2:], mode='bilinear', align_corners=True) x = torch.cat([x, low_level_feat], dim=1) x = self.last_conv(x) return x

深度可分离卷积的实现及其与普通卷积的对比：

# 普通卷积 class RegularConv(nn.Module): def __init__(self, in_channels, out_channels, kernel_size=3): super().__init__() self.conv = nn.Conv2d( in_channels, out_channels, kernel_size, padding=kernel_size//2, bias=False ) self.bn = nn.BatchNorm2d(out_channels) def forward(self, x): return F.relu(self.bn(self.conv(x))) # 深度可分离卷积 class SeparableConv(nn.Module): def __init__(self, in_channels, out_channels, kernel_size=3): super().__init__() self.depthwise = nn.Conv2d( in_channels, in_channels, kernel_size, padding=kernel_size//2, groups=in_channels, bias=False ) self.pointwise = nn.Conv2d(in_channels, out_channels, 1, bias=False) self.bn = nn.BatchNorm2d(out_channels) def forward(self, x): x = self.depthwise(x) x = self.pointwise(x) return F.relu(self.bn(x))

两种卷积的参数数量对比（假设in_channels=256, out_channels=256, kernel_size=3）：

在实际项目中，将ASPP中的常规卷积替换为深度可分离卷积可以显著减少计算量：

class AtrousSeparableConv(nn.Module): def __init__(self, in_channels, out_channels, dilation): super().__init__() self.depthwise = nn.Conv2d( in_channels, in_channels, 3, padding=dilation, dilation=dilation, groups=in_channels, bias=False ) self.pointwise = nn.Conv2d(in_channels, out_channels, 1, bias=False) self.bn = nn.BatchNorm2d(out_channels) def forward(self, x): x = self.depthwise(x) x = self.pointwise(x) return F.relu(self.bn(x))

6. 完整模型集成与训练技巧

将上述模块组合成完整的DeepLabV3+模型：

class DeepLabV3Plus(nn.Module): def __init__(self, backbone='resnet50', num_classes=21, output_stride=16): super().__init__() # 根据output_stride设置dilation rates if output_stride == 16: rates = [1, 6, 12, 18] aspp_rates = [6, 12, 18] else: # output_stride=8 rates = [1, 12, 24, 36] aspp_rates = [12, 24, 36] # 构建骨干网络 self.backbone = build_backbone(backbone, output_stride) low_level_channels = self.backbone.low_level_channels # ASPP模块 self.aspp = ASPP(self.backbone.out_channels, 256, aspp_rates) # 解码器 self.decoder = Decoder(low_level_channels, num_classes) # 初始化权重 self._init_weight() def forward(self, x): size = x.shape[2:] # 编码器部分 x, low_level_feat = self.backbone(x) # ASPP部分 x = self.aspp(x) # 解码器部分 x = self.decoder(x, low_level_feat) # 上采样到原图大小 x = F.interpolate(x, size=size, mode='bilinear', align_corners=True) return x def _init_weight(self): for m in self.modules(): if isinstance(m, nn.Conv2d): nn.init.kaiming_normal_(m.weight) elif isinstance(m, nn.BatchNorm2d): m.weight.data.fill_(1) m.bias.data.zero_()

训练DeepLab模型时需要注意的关键点：

1. 学习率策略：

   • 使用多项式学习率衰减：$lr = base_lr \times (1 - \frac{iter}{max_iter})^{power}$
   • 典型设置：base_lr=0.007, power=0.9

2. 数据增强：

   • 随机缩放（0.5-2.0倍）
   • 随机左右翻转
   • 随机裁剪（通常为513×513）

3. 损失函数：

   • 交叉熵损失为主损失
   • 可辅助使用辅助损失（auxiliary loss）

def create_optimizer(model, base_lr=0.007, momentum=0.9, weight_decay=0.0005): params_dict = dict(model.named_parameters()) params = [] for key, value in params_dict.items(): if 'backbone' in key: params += [{'params': [value], 'lr': base_lr * 0.1}] else: params += [{'params': [value], 'lr': base_lr}] optimizer = torch.optim.SGD(params, momentum=momentum, weight_decay=weight_decay) return optimizer

在Cityscapes数据集上的典型训练配置：