AOD-Net预训练权重实战指南：从加载到微调的全流程解析

1次阅读

没有评论

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

AOD-Net（All-in-One Dehazing Network）是端到端的轻量级去雾网络，相比传统基于物理模型的方法有三大优势：

直接学习雾图到清晰图的映射，避免透射率 / 大气光估计误差累积
单阶段网络结构在移动端部署时推理速度可达 25FPS
预训练权重在 RESIDE 数据集上 PSNR 达 23.5，优于多数传统方法

根据社区反馈，初次使用预训练权重时 90% 的问题集中在：

环境配置问题
PyTorch 与 CUDA 版本不匹配引发 undefined symbol: cudaSetupArgument 错误
张量维度错误
输入图像未做归一化或尺寸不符导致RuntimeError: size mismatch
权重加载失败
直接加载 .pth 文件出现 Missing key(s) in state_dict 警告

推荐从官方 GitHub 仓库下载标准权重：

import requests

# 官方权重下载链接（建议国内用户配置代理）url = 'https://github.com/MayankSingal/AOD-Net-PyTorch/releases/download/v1.0/AOD_net_weights.pth'
r = requests.get(url, allow_redirects=True)
open('AOD_weights.pth', 'wb').write(r.content)

使用 PyTorch 的严格加载模式避免参数错位：

import torch
from model import AOD_net  # 需提前定义网络结构

def load_pretrained(model, weight_path):
    try:
        # 建议始终指定 map_location 避免跨设备问题
        state_dict = torch.load(weight_path, map_location='cuda:0' if torch.cuda.is_available() else 'cpu')

        # 过滤掉可能存在的 module. 前缀（多卡训练保存的权重）new_state_dict = {k.replace('module.', ''): v for k, v in state_dict.items()}

        # 严格匹配模型结构
        model.load_state_dict(new_state_dict, strict=True)
        print('✅ 权重加载成功')
        return model
    except Exception as e:
        print(f'❌ 加载失败: {str(e)}')
        # 失败时返回未经训练的模型
        return model

net = AOD_net()
net = load_pretrained(net, 'AOD_weights.pth')

AOD-Net 输入需要满足：

图像缩放至固定尺寸（原论文使用 256×256）
像素值归一化到 [-1, 1] 范围

from torchvision import transforms

train_transform = transforms.Compose([transforms.Resize((256, 256)),  # 建议保持与预训练相同尺寸
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])  # 归一化到[-1,1]
])

推荐采用分阶段学习率：

前 5 个 epoch：使用 1e- 4 的小学习率微调最后一层
后续训练：增大到 5e- 4 解冻全部层

import torch.optim as optim

# 分组设置学习率
optimizer = optim.Adam([{'params': net.dehaze[-1].parameters(), 'lr': 1e-4},  # 最后一层
    {'params': net.encoder.parameters(), 'lr': 5e-5}     # 其余层
])

# 添加学习率预热
scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=0.1, total_iters=5)

当出现 CUDA out of memory 时：

梯度累积：每 4 个 batch 更新一次参数

for i, batch in enumerate(dataloader):
    loss = model(batch)
    loss = loss / 4  # 梯度累加
    loss.backward()

    if (i+1) % 4 == 0:
        optimizer.step()
        optimizer.zero_grad()

混合精度训练：减少显存占用约 40%

from torch.cuda.amp import autocast, GradScaler

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

在 SOTS 户外测试集上的对比数据：