从ResNet到ASPP：手把手教你用PyTorch复现DeepLabv3+的Encoder模块（含代码详解）
·
从ResNet到ASPP：手把手教你用PyTorch复现DeepLabv3+的Encoder模块（含代码详解）
在语义分割领域，DeepLabv3+以其出色的性能和清晰的架构设计成为众多研究者和工程师的首选方案。本文将带您深入探索其核心组件——Encoder模块的实现细节，从ResNet-101骨干网络到ASPP（Atrous Spatial Pyramid Pooling）结构，通过PyTorch代码逐行解析，帮助您彻底掌握这一关键技术。
1. 环境准备与基础架构
在开始编码之前，我们需要搭建好开发环境并理解DeepLabv3+ Encoder的整体架构。以下是推荐的环境配置：
# 环境配置要求
import torch
import torch.nn as nn
import torchvision
print(f"PyTorch版本: {torch.__version__}")
print(f"Torchvision版本: {torchvision.__version__}")
print(f"CUDA可用: {torch.cuda.is_available()}")
DeepLabv3+的Encoder由两部分组成：
- 骨干网络(Backbone): 通常采用ResNet-101提取多层次特征
- ASPP模块: 通过不同膨胀率的空洞卷积捕获多尺度上下文信息
提示：建议使用Python 3.8+和PyTorch 1.10+版本以获得最佳兼容性
2. ResNet-101骨干网络实现
ResNet作为Encoder的核心组件，其实现需要特别注意空洞卷积的改造。我们将基于torchvision的预训练模型进行修改：
class ResNetBackbone(nn.Module):
def __init__(self, pretrained=True):
super().__init__()
# 加载预训练ResNet-101
resnet = torchvision.models.resnet101(pretrained=pretrained)
# 提取各阶段特征提取层
self.conv1 = resnet.conv1
self.bn1 = resnet.bn1
self.relu = resnet.relu
self.maxpool = resnet.maxpool
self.layer1 = resnet.layer1 # 输出stride=4
self.layer2 = resnet.layer2 # 输出stride=8
self.layer3 = resnet.layer3 # 输出stride=16
self.layer4 = resnet.layer4 # 输出stride=32
# 将layer3和layer4的stride从2改为1
self._modify_stride(self.layer3)
self._modify_stride(self.layer4)
# 为layer3和layer4添加空洞卷积
self._apply_dilation(self.layer3, dilation=2)
self._apply_dilation(self.layer4, dilation=4)
def _modify_stride(self, layer):
"""将指定层的stride从2改为1"""
for block in layer:
if isinstance(block, torchvision.models.resnet.Bottleneck):
if block.downsample is not None:
block.downsample[0].stride = (1, 1)
block.conv2.stride = (1, 1)
def _apply_dilation(self, layer, dilation):
"""为指定层添加空洞卷积"""
for block in layer:
if isinstance(block, torchvision.models.resnet.Bottleneck):
block.conv2.dilation = (dilation, dilation)
block.conv2.padding = (dilation, dilation)
def forward(self, x):
# 前向传播过程
x = self.conv1(x)
x = self.bn1(x)
x = self.relu(x)
x = self.maxpool(x)
x = self.layer1(x) # stride=4
low_level_feat = x # 保存低级特征供Decoder使用
x = self.layer2(x) # stride=8
x = self.layer3(x) # stride=16 (修改后)
x = self.layer4(x) # stride=16 (修改后)
return x, low_level_feat
关键修改点说明：
- stride调整：将layer3和layer4的stride从2改为1，避免特征图过度缩小
- 空洞卷积应用：为layer3和layer4添加dilation参数，扩大感受野
- 特征保留：保存layer1输出的低级特征(low_level_feat)供Decoder使用
3. ASPP模块实现详解
ASPP模块是DeepLabv3+的核心创新，它通过并行使用不同膨胀率的空洞卷积捕获多尺度信息：
class ASPP(nn.Module):
def __init__(self, in_channels, out_channels=256, atrous_rates=[6, 12, 18]):
super().__init__()
# 1x1卷积分支
self.conv1x1 = nn.Sequential(
nn.Conv2d(in_channels, out_channels, 1, bias=False),
nn.BatchNorm2d(out_channels),
nn.ReLU()
)
# 3x3卷积分支（不同膨胀率）
self.conv3x3_1 = self._make_aspp_conv(in_channels, out_channels, atrous_rates[0])
self.conv3x3_2 = self._make_aspp_conv(in_channels, out_channels, atrous_rates[1])
self.conv3x3_3 = self._make_aspp_conv(in_channels, out_channels, atrous_rates[2])
# 图像级特征分支（全局平均池化+1x1卷积）
self.image_pool = nn.Sequential(
nn.AdaptiveAvgPool2d(1),
nn.Conv2d(in_channels, out_channels, 1, bias=False),
nn.BatchNorm2d(out_channels),
nn.ReLU()
)
# 输出卷积层
self.conv_out = nn.Sequential(
nn.Conv2d(out_channels*5, out_channels, 1, bias=False),
nn.BatchNorm2d(out_channels),
nn.ReLU(),
nn.Dropout(0.5)
)
def _make_aspp_conv(self, in_channels, out_channels, dilation):
return nn.Sequential(
nn.Conv2d(in_channels, out_channels, 3,
padding=dilation, dilation=dilation, bias=False),
nn.BatchNorm2d(out_channels),
nn.ReLU()
)
def forward(self, x):
# 获取输入特征图尺寸
h, w = x.size()[2:]
# 各分支处理
conv1x1 = self.conv1x1(x)
conv3x3_1 = self.conv3x3_1(x)
conv3x3_2 = self.conv3x3_2(x)
conv3x3_3 = self.conv3x3_3(x)
# 图像级特征处理
img_feat = self.image_pool(x)
img_feat = F.interpolate(img_feat, size=(h, w),
mode='bilinear', align_corners=True)
# 特征拼接
x = torch.cat([conv1x1, conv3x3_1, conv3x3_2, conv3x3_3, img_feat], dim=1)
x = self.conv_out(x)
return x
ASPP模块包含五个并行分支：
- 1x1卷积：捕获局部特征
- 3x3空洞卷积(rate=6)：中等感受野
- 3x3空洞卷积(rate=12)：较大感受野
- 3x3空洞卷积(rate=18)：最大感受野
- 图像级特征：全局上下文信息
4. Encoder模块完整实现与测试
将ResNet骨干网络和ASPP模块组合成完整的Encoder：
class DeepLabV3PlusEncoder(nn.Module):
def __init__(self, num_classes=21, pretrained=True):
super().__init__()
# 骨干网络
self.backbone = ResNetBackbone(pretrained=pretrained)
# ASPP模块
self.aspp = ASPP(in_channels=2048) # ResNet-101最后一层通道数为2048
# 低级特征处理
self.low_level_conv = nn.Sequential(
nn.Conv2d(256, 48, 1, bias=False), # ResNet layer1输出256通道
nn.BatchNorm2d(48),
nn.ReLU()
)
# 分类头（实际使用时Decoder会替换这部分）
self.classifier = nn.Sequential(
nn.Conv2d(304, 256, 3, padding=1, bias=False), # 256+48=304
nn.BatchNorm2d(256),
nn.ReLU(),
nn.Conv2d(256, num_classes, 1)
)
def forward(self, x):
# 骨干网络前向传播
x, low_level_feat = self.backbone(x)
# ASPP处理高级特征
x = self.aspp(x)
x = F.interpolate(x, scale_factor=4, mode='bilinear', align_corners=True)
# 处理低级特征
low_level_feat = self.low_level_conv(low_level_feat)
# 特征融合
x = torch.cat([x, low_level_feat], dim=1)
x = self.classifier(x)
x = F.interpolate(x, scale_factor=4, mode='bilinear', align_corners=True)
return x
测试Encoder的完整流程：
# 测试代码
if __name__ == '__main__':
# 创建模型实例
model = DeepLabV3PlusEncoder(num_classes=21)
# 模拟输入 (batch_size=1, channels=3, height=512, width=512)
dummy_input = torch.randn(1, 3, 512, 512)
# 前向传播
output = model(dummy_input)
print(f"输入尺寸: {dummy_input.shape}")
print(f"输出尺寸: {output.shape}") # 应为(1, 21, 512, 512)
5. 关键问题与解决方案
在实际实现过程中，我们可能会遇到以下几个典型问题：
5.1 特征图尺寸对齐问题
当融合不同层次的特征时，尺寸不匹配是常见问题。我们的解决方案包括：
-
精确计算各层输出尺寸：使用以下公式计算空洞卷积后的特征图大小：
H_out = floor[(H_in + 2*padding - dilation*(kernel_size-1) -1)/stride +1] -
使用双线性插值调整尺寸：在特征融合前统一尺寸
5.2 内存消耗优化
DeepLabv3+的Encoder可能消耗大量显存，特别是处理高分辨率图像时。优化策略：
-
梯度检查点技术：
from torch.utils.checkpoint import checkpoint # 在forward方法中使用 x = checkpoint(self.layer3, x) -
混合精度训练：
from torch.cuda.amp import autocast with autocast(): output = model(input)
5.3 训练技巧与参数调优
基于实际项目经验，推荐以下训练配置：
| 参数 | 推荐值 | 说明 |
|---|---|---|
| 学习率 | 0.007 | 使用poly学习率衰减策略 |
| 批量大小 | 16 | 根据GPU显存调整 |
| 优化器 | SGD | momentum=0.9, weight_decay=0.0005 |
| 训练epoch | 50 | 在Cityscapes等大数据集上 |
注意：当使用预训练模型时，建议骨干网络采用较小的学习率（如主学习率的0.1倍）
6. 性能评估与可视化
为了验证Encoder的实现正确性，我们可以进行以下测试：
-
感受野可视化：
def visualize_receptive_field(model, input_size=(512, 512)): from torchvision.models.feature_extraction import create_feature_extractor # 创建特征提取器 model.eval() nodes = {'aspp.conv3x3_3.0': 'output'} extractor = create_feature_extractor(model, return_nodes=nodes) # 生成测试图像 img = torch.zeros(1, 3, *input_size) center = (input_size[0]//2, input_size[1]//2) img[0, :, center[0], center[1]] = 1 # 计算梯度 img.requires_grad = True output = extractor(img)['output'] output.sum().backward() # 可视化梯度 grad_img = img.grad[0].sum(dim=0).detach().numpy() plt.imshow(grad_img, cmap='hot') plt.title('Receptive Field') -
特征图可视化：
def visualize_features(model, input_image): # 获取各层特征 features = {} def hook_fn(name): def hook(module, input, output): features[name] = output.detach() return hook hooks = [] for name, layer in model.named_children(): hooks.append(layer.register_forward_hook(hook_fn(name))) # 前向传播 with torch.no_grad(): _ = model(input_image) # 移除钩子 for hook in hooks: hook.remove() # 可视化特征 fig, axes = plt.subplots(2, 3, figsize=(15, 10)) for i, (name, feat) in enumerate(features.items()): ax = axes[i//3, i%3] ax.imshow(feat[0, 0].cpu().numpy(), cmap='viridis') ax.set_title(name)
7. 高级优化技巧
对于追求更高性能的开发者，可以考虑以下进阶优化：
-
可变形卷积替代空洞卷积：
from torchvision.ops import DeformConv2d class DeformableASPPConv(nn.Module): def __init__(self, in_channels, out_channels, dilation): super().__init__() self.offset = nn.Conv2d(in_channels, 2*3*3, 3, padding=dilation, dilation=dilation) self.conv = DeformConv2d(in_channels, out_channels, 3, padding=dilation, dilation=dilation) def forward(self, x): offset = self.offset(x) return self.conv(x, offset) -
注意力机制增强：
class ChannelAttention(nn.Module): def __init__(self, in_channels, ratio=8): super().__init__() self.avg_pool = nn.AdaptiveAvgPool2d(1) self.max_pool = nn.AdaptiveMaxPool2d(1) self.fc = nn.Sequential( nn.Linear(in_channels, in_channels//ratio), nn.ReLU(), nn.Linear(in_channels//ratio, in_channels) ) def forward(self, x): b, c, _, _ = x.size() avg_out = self.fc(self.avg_pool(x).view(b, c)) max_out = self.fc(self.max_pool(x).view(b, c)) out = avg_out + max_out return torch.sigmoid(out).view(b, c, 1, 1) * x -
知识蒸馏压缩模型：
class DistillationLoss(nn.Module): def __init__(self, T=2.0): super().__init__() self.T = T self.kl_div = nn.KLDivLoss(reduction='batchmean') def forward(self, student_out, teacher_out): soft_student = F.log_softmax(student_out/self.T, dim=1) soft_teacher = F.softmax(teacher_out/self.T, dim=1) return self.kl_div(soft_student, soft_teacher) * (self.T**2)
在实际项目中，这些优化技巧可以将模型mIoU提升2-5个百分点，但也会相应增加实现复杂度。建议先完成基础版本，再逐步引入高级优化。
码道开发者社区，聚焦华为云码道 CodeArts 代码智能体，沉淀 Agent、Skill、鸿蒙开发实战内容，供开发者查阅资料、交流技术、分享工程实践
更多推荐
1
0- 0
已为社区贡献116条内容
所有评论(0)