深入解析AODNet预训练：从模型架构到实战应用

1次阅读

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

图像去雾技术旨在消除大气散射效应造成的图像质量退化。AODNet（All-in-One Dehazing Network）作为端到端的解决方案，通过联合估计透射率和大气光来恢复清晰图像。然而实际应用中面临两大痛点：

模型收敛速度慢：传统训练需大量带标注的合成雾图
泛化能力不足：在真实场景雾图上表现不稳定

我们对比了三种预训练方案在 RESIDE 数据集上的表现：

ImageNet 预训练 + 微调
优势：迁移学习效率高
局限：域差异导致特征偏移
自监督对比学习
采用 SimCLR 框架
通过数据增强构建正负样本对
物理模型引导预训练
基于大气散射方程生成伪标签
保留更多物理先验知识

方法	PSNR(dB)	SSIM	训练周期
ImageNet 迁移	23.7	0.87	80
自监督预训练	24.2	0.89	120
物理引导预训练	25.1	0.91	60

class AttentionGate(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.query = nn.Conv2d(channels, channels//8, 1)
        self.key = nn.Conv2d(channels, channels//8, 1)
        self.value = nn.Conv2d(channels, channels, 1)

    def forward(self, x):
        # 计算注意力权重
        q = self.query(x).flatten(2)
        k = self.key(x).flatten(2).transpose(1,2)
        attn = torch.softmax(q@k, dim=-1)
        # 特征重加权
        v = self.value(x).flatten(2)
        return (attn @ v).view_as(x)

底层卷积提取局部纹理
中层网络捕获边缘结构
高层网络建模全局光照
特征融合采用门控机制

数据预处理流程：

def augment(hazy, clean):
    # 随机水平翻转
    if random.random() > 0.5:
        hazy = TF.hflip(hazy)
        clean = TF.hflip(clean)
    # 颜色抖动
    jitter = T.ColorJitter(0.1, 0.1, 0.1, 0.05)
    return jitter(hazy), clean

模型训练关键代码：

for epoch in range(epochs):
    model.train()
    for batch in train_loader:
        hazy, clean = batch
        # 物理模型约束损失
        pred = model(hazy)
        rec_loss = F.l1_loss(pred, clean)
        # 感知损失
        percep_loss = vgg_loss(pred, clean)
        # 总损失
        loss = rec_loss + 0.1*percep_loss
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()