共计 2698 个字符，预计需要花费 7 分钟才能阅读完成。
背景痛点
高分辨率图像的语义分割任务对细节保留和长距离依赖建模提出了极高要求。传统 CNN 方法在这方面存在明显不足：
- 感受野限制 ：CNN 的局部感受野难以捕捉图像中相距较远的像素间关系，导致长距离依赖建模能力弱。
- 细节丢失 ：多次下采样操作会丢失细粒度空间信息，影响边界分割精度。
- 固定权重 ：卷积核权重在推理时固定，难以自适应不同输入内容。
Transformer 的自注意力机制恰好能解决这些问题：
- 全局建模 ：自注意力可以建立任意像素间的关系，不受距离限制。
- 内容自适应 ：attention 权重动态计算，能更好适应不同图像区域的特征。
- 细节保留 ：通过合适的 patch 划分，可以保持高分辨率特征。
方案对比
我们对比了几种主流分割方案在 Cityscapes 数据集上的表现：
| 方法类型 | 代表模型 | mIoU(%) | 推理速度 (FPS) |
|---|---|---|---|
| CNN-based | DeepLabV3+ | 79.3 | 8.7 |
| CNN-based | UNet | 76.9 | 12.1 |
| Transformer | Segmenter | 81.2 | 6.3 |
| Transformer | MaskFormer | 83.5 | 4.8 |
可以看到 Transformer 方法在精度上有明显优势，但速度较慢，这也是后续优化的重点方向。
核心实现
Patch Embedding 实现
import torch
import torch.nn as nn
class PatchEmbed(nn.Module):
"""
将图像分割为 patch 并嵌入到特征空间
Args:
img_size (int): 输入图像尺寸
patch_size (int): patch 大小
in_chans (int): 输入通道数
embed_dim (int): 嵌入维度
overlap (float): patch 重叠率 (0-0.5)
"""
def __init__(self, img_size=224, patch_size=16, in_chans=3,
embed_dim=768, overlap=0.25):
super().__init__()
# 计算实际步长 (考虑重叠)
stride = int(patch_size * (1 - overlap))
# 重叠卷积实现
self.proj = nn.Conv2d(in_chans, embed_dim,
kernel_size=patch_size,
stride=stride,
padding=patch_size//2)
# 标准化
self.norm = nn.LayerNorm(embed_dim)
def forward(self, x):
x = self.proj(x) # (B, C, H, W) -> (B, D, H', W')
x = x.flatten(2).transpose(1, 2) # (B, D, H'*W') -> (B, H'*W', D)
x = self.norm(x)
return x 关键参数说明：
– overlap 建议 0.25-0.4 之间，过大增加计算量，过小丢失信息
– embed_dim 通常设置为 768 或 1024
解码器优化
解码器采用经典的 UNet 结构，但加入了以下改进：
- 多尺度特征融合 ：在跳跃连接中加入特征选择模块
- 注意力引导 ：在高层特征和低层特征 concat 前加入 cross-attention
- 渐进上采样 ：采用 2 倍逐步上采样而非直接恢复到原图尺寸
性能优化
混合精度训练
from torch.cuda.amp import autocast, GradScaler
scaler = GradScaler() # 防止梯度下溢
for inputs, labels in dataloader:
optimizer.zero_grad()
with autocast():
outputs = model(inputs)
loss = criterion(outputs, labels)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update() 配置建议：
– 初始 scale 值设为 65536.0
– 每 200 次迭代检查是否发生 inf/nan
梯度检查点
from torch.utils.checkpoint import checkpoint
class TransformerBlock(nn.Module):
def forward(self, x):
return checkpoint(self._forward, x)
def _forward(self, x):
# 正常的 transformer 前向计算
return x可节省约 60% 显存，但训练速度会降低 20-30%。
避坑指南
位置编码适配
高分辨率图像下位置编码需要特别注意：
- 绝对位置编码需要调整最大长度
- 相对位置偏置的窗口大小需增大
- 建议使用可学习的位置编码而非固定公式
小样本预训练
数据不足时建议：
- 使用 ImageNet-21k 预训练的 ViT 作为 backbone
- 冻结前几层 transformer block
- 采用 heavy augmentation
延伸思考
一个值得探讨的问题是：如何平衡注意力计算复杂度与分割精度？
当前方案的计算复杂度是 O(N²)，对于高分辨率图像压力较大。可能的改进方向：
- 窗口注意力 (Swin Transformer 方案)
- 轴向注意力 (Axial-DeepLab 方案)
- 稀疏注意力 (如 Longformer)
读者可以尝试在现有代码基础上实现窗口注意力：
# 示例：窗口注意力实现
window_size = 8
B, L, C = x.shape
H = W = int(L**0.5)
x = x.view(B, H, W, C)
# 划分窗口
x = x.reshape(B, H//window_size, window_size, W//window_size, window_size, C)
x = x.permute(0,1,3,2,4,5).reshape(-1, window_size*window_size, C)
# 窗口内做自注意力
attn = (x @ x.transpose(-2,-1)) / (C**0.5)
attn = attn.softmax(dim=-1)
x = attn @ x
# 恢复原始形状
x = x.view(B, H//window_size, W//window_size, window_size, window_size, C)
x = x.permute(0,1,3,2,4,5).reshape(B, H, W, C)
x = x.reshape(B, L, C)这种方法可以将复杂度降至 O(N)，是很好的优化方向。
结语
Transformer 在语义分割领域展现出强大潜力，特别是在需要精细分割的高分辨率场景。通过合理的架构设计和优化技巧，我们可以在保持精度的前提下提升推理效率。希望本文的实践经验和优化建议能帮助读者在自己的项目中更好地应用这一技术。
正文完
