基于Attention Retractable Transformer的图像精准修复：从原理到新手实践

共计 2015 个字符，预计需要花费 6 分钟才能阅读完成。

背景痛点：为什么需要新架构？

传统 CNN 在图像修复中存在三大局限性：

1. 感受野受限 ：卷积核尺寸固定，难以建模长距离依赖关系。修复大范围缺失区域时，容易出现结构扭曲（如脸部的对称性破坏）
2. 内容模糊 ：反复下采样 - 上采样过程中丢失高频细节，导致修复区域出现明显模糊块效应
3. 计算冗余 ：对每个像素使用相同卷积核，无法针对不同区域复杂度动态调整计算量

普通 Transformer 虽能解决长程依赖问题，但带来新挑战：

• 复杂度爆炸 ：标准自注意力计算复杂度为 O(N²)，处理 512×512 图像时内存占用超过 32GB
• 局部细节丢失 ：全局注意力使模型过度关注显著区域，忽略细粒度纹理（如发丝、皮肤毛孔）

技术对比：架构进化之路

测试环境：CelebA-HQ 256×256，RTX 3090

核心实现：可伸缩注意力机制

动态范围计算（数学推导）

定义注意力范围为可学习参数：

r = \sigma(W_r \cdot \text{AvgPool}(F_{in})) \times R_{max}

其中 $R_{max}$ 为预设最大范围，$\sigma$ 为 sigmoid 函数。代码实现：

class DynamicRange(nn.Module):
    def __init__(self, dim, R_max=32):
        super().__init__()
        self.range_pred = nn.Sequential(nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(dim, dim//4, 1),
            nn.ReLU(),
            nn.Conv2d(dim//4, 1, 1),
            nn.Sigmoid())
        self.R_max = R_max

    def forward(self, x):
        return self.range_pred(x) * self.R_max  # [B,1,1,1]

多尺度特征融合

（示意图：金字塔结构包含 4 个尺度，每层注意力范围递减）

关键实现步骤：

1. 下采样生成特征金字塔 {F1,F2,F3,F4}
2. 每层计算独立注意力范围 r_i = R_max/2^(i-1)
3. 跨尺度信息通过门控机制融合：

   # 门控权重计算
   gate = torch.sigmoid(self.gate_conv(torch.cat([low_feat, high_feat], dim=1)))
   fused_feat = gate * high_feat + (1-gate) * F.interpolate(low_feat, scale_factor=2)

实验验证：CelebA-HQ 结果

注：测试集包含 3000 张遮挡率 30%~50% 的图像

避坑指南：实战经验

学习率策略

• 初始阶段 (0-10k iter)：固定 lr=1e-4
• 中期 (10k-50k)：余弦退火到 5e-5
• 后期 (>50k)：启用梯度裁剪 (thresh=0.1)

显存优化三连

1. 使用混合精度训练（AMP）

   scaler = torch.cuda.amp.GradScaler()
   with autocast():
       pred = model(x)
       loss = criterion(pred, y)
   scaler.scale(loss).backward()
   scaler.step(optimizer)

2. 梯度累积：每 4 个 batch 更新一次
3. 激活检查点：对 Transformer 层启用

   model.set_grad_checkpointing(True)  # 节省 30% 显存 

注意力坍塌诊断

症状：修复区域出现重复纹理（如棋盘格）
解决方案：

1. 监控注意力熵：$H = -\sum p_i\log p_i$ 应大于 2.5
2. 添加多样性损失：

   L_{div} = \frac{1}{N}\sum_{i=1}^N \max(0, \cos(A_i,A_j)-0.5)

延伸思考：视频修复迁移

时序扩展方案：

1. 3D 注意力窗口 ：在空间维度保持 retractable 特性，时间维度固定较小范围（如 5 帧）
2. 光流引导 ：利用相邻帧运动信息初始化注意力中心位置
3. 缓存机制 ：对静态背景区域复用前一帧特征，减少 60% 计算量

核心代码修改点：

# 将 2D 卷积替换为伪 3D 卷积
self.temp_conv = nn.Conv3d(in_dim, out_dim, kernel_size=(3,1,1), padding=(1,0,0))

结语

通过可伸缩注意力机制，ART 在精度和效率间取得巧妙平衡。建议初学者先从 CelebA-HQ 小尺度图像开始实验，逐步掌握动态范围调整的技巧。遇到训练波动时，可尝试冻结注意力模块先训练编码器部分。期待看到大家创造出更有创意的变体！