AODNet预训练实战指南：从零搭建到性能调优

1次阅读

没有评论

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

刚接触 AODNet 预训练时，最容易遇到两个拦路虎：

梯度消失问题：当网络深度增加时，靠近输入层的参数更新幅度变得极小。在 AODNet 中表现为去雾效果随训练轮次提升缓慢
训练不稳定：损失函数值会出现周期性波动，尤其在 batch size 较大时（如 >32），这种现象在 Adam 优化器中更为明显

通过实验发现，这些问题与初始化策略和优化器选择密切相关。例如使用 Xavier 初始化时，第一层卷积的梯度范数比最后一层小 3 个数量级。

在 RTX3090+PyTorch1.12 环境下测试了两种优化器：

AdamW (weight decay=0.05)
优点：默认参数下收敛稳定
缺点：最终 PSNR 比 LAMB 低约 0.8dB
适用场景：小规模数据（<10 万样本）
LAMB (layer-wise adaptive rate)
优点：大 batch(256)下仍保持稳定
缺点：需要预热 (warmup) 阶段
内存开销：比 AdamW 多 15%

实测数据对比表：

优化器	最终 PSNR	训练时长	显存占用
AdamW	28.7dB	4.2h	9.8GB
LAMB	29.5dB	3.8h	11.3GB

使用 Albumentations 构建预处理流水线：

import albumentations as A

train_transform = A.Compose([A.RandomRotate90(),  # 随机旋转
    A.Flip(p=0.5),  # 水平翻转
    A.RandomBrightnessContrast(p=0.2),  # 亮度对比度调整
    A.GaussNoise(var_limit=(0, 0.01)),  # 高斯噪声
    A.Normalize(mean=[0.485, 0.456, 0.406], 
                std=[0.229, 0.224, 0.225])  # ImageNet 归一化
], additional_targets={'depth': 'image'})  # 支持深度图输入

梯度裁剪：限制最大梯度范数为 1.0

torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)

EMA(指数移动平均)：

class ModelEMA:
    def __init__(self, model, decay=0.9999):
        self.ema = deepcopy(model).eval()
        self.decay = decay

    def update(self, model):
        with torch.no_grad():
            for ema_p, model_p in zip(self.ema.parameters(), model.parameters()):
                ema_p.mul_(self.decay).add_(model_p, alpha=1-self.decay)

启用 AMP 自动混合精度后，显存占用从 12.1GB 降至 8.4GB：

scaler = torch.cuda.amp.GradScaler()

with autocast():
    output = model(input)
    loss = criterion(output, target)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

使用 TorchProfiler 定位到三个耗时操作：

转置卷积 (ConvTranspose2d) 占总时长 35%
特征图拼接 (concat) 占 18%
激活函数 (SiLU) 占 12%

优化建议：
– 将转置卷积替换为最近邻上采样 + 普通卷积
– 预分配拼接后的张量内存

导致验证集指标震荡的典型错误：

学习率与 batch size 不匹配 ：当 batch size 扩大 N 倍时，学习率应同步扩大 sqrt(N) 倍
错误的数据标准化：对 HDR 图像错误使用 0 - 1 归一化会导致数值溢出
验证集数据泄露：在训练流程中意外对验证集执行了增强操作

关键卷积操作应标注张量形状：

# [batch, channels, height, width]
x = torch.randn(4, 3, 256, 256)  # 输入张量
conv = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1)
y = conv(x)  # 输出形状 [4, 64, 256, 256]

关于 AODNet 的两个未解之谜：