从零构建HAN超分网络：实战通道与空间注意力机制

在计算机视觉领域，单图像超分辨率(SISR)一直是备受关注的研究方向。随着深度学习技术的发展，基于注意力机制的神经网络在超分任务中展现出显著优势。2020年ECCV会议上提出的Holistic Attention Network(HAN)通过创新的层注意力模块(LAM)和通道空间注意力模块(CSAM)，在多个基准测试集上达到了当时最先进的性能。本文将带您从零开始实现这个网络，不仅解析其核心设计思想，还会分享实际训练中的调参技巧。

1. 环境准备与数据加载

构建HAN网络的第一步是搭建合适的开发环境。推荐使用Python 3.8+和PyTorch 1.7+的组合，这些版本在兼容性和性能方面都经过了充分验证。以下是基础环境配置步骤：

bash复制conda create -n han python=3.8
conda activate han
pip install torch==1.7.1 torchvision==0.8.2
pip install opencv-python numpy tqdm

对于训练数据，DIV2K数据集是超分任务的标准选择。这个数据集包含800张训练图像和100张验证图像，每张都有对应的高分辨率(HR)和低分辨率(LR)版本。在实际应用中，我们通常需要模拟不同的降质过程：

python复制from torchvision import transforms

class DegradationProcess:
    def __init__(self, scale_factor=4, blur_kernel=7):
        self.scale_factor = scale_factor
        self.blur_kernel = blur_kernel
        
    def __call__(self, hr_img):
        # 双三次下采样
        lr_img = transforms.Resize(
            hr_img.size[1]//self.scale_factor,
            interpolation=transforms.InterpolationMode.BICUBIC
        )(hr_img)
        # 添加高斯模糊
        lr_img = transforms.GaussianBlur(self.blur_kernel)(lr_img)
        return lr_img

提示：在实际项目中，建议将预处理后的数据保存为.npy文件，可以显著减少训练时的IO开销。

2. 网络架构设计与核心模块

HAN网络的核心创新在于其注意力机制的设计，它突破了传统单层优化的局限，实现了全局特征协作。整体架构包含四个关键部分：浅层特征提取、残差组堆叠、层注意力模块和通道空间注意力模块。

2.1 基础特征提取模块

网络的第一阶段是浅层特征提取，这部分的实现相对简单但非常重要：

python复制import torch.nn as nn

class ShallowFeatureExtraction(nn.Module):
    def __init__(self, in_channels=3, out_channels=64):
        super().__init__()
        self.conv = nn.Conv2d(
            in_channels, out_channels,
            kernel_size=3, padding=1
        )
        
    def forward(self, x):
        return self.conv(x)

2.2 残差组与RCAN结构

HAN采用了类似RCAN的残差组设计，每个组包含多个残差通道注意力块：

python复制class RCAB(nn.Module):
    def __init__(self, channels=64, reduction=16):
        super().__init__()
        self.body = nn.Sequential(
            nn.Conv2d(channels, channels, 3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(channels, channels, 3, padding=1),
            ChannelAttention(channels, reduction)
        )
        
    def forward(self, x):
        return x + self.body(x)

class ResidualGroup(nn.Module):
    def __init__(self, n_blocks=20, channels=64):
        super().__init__()
        self.blocks = nn.Sequential(
            *[RCAB(channels) for _ in range(n_blocks)]
        )
        self.conv = nn.Conv2d(channels, channels, 3, padding=1)
        
    def forward(self, x):
        return x + self.conv(self.blocks(x))

3. 注意力机制实现

HAN的核心创新在于其注意力机制设计，下面我们详细实现这两个关键模块。

3.1 层注意力模块(LAM)

LAM模块通过建模不同层特征间的关系，实现了全局层面的特征优化：

python复制class LAM(nn.Module):
    def __init__(self, n_resgroups=10, channels=64):
        super().__init__()
        self.n_resgroups = n_resgroups
        self.channels = channels
        self.alpha = nn.Parameter(torch.zeros(1))
        
    def forward(self, features):
        # features: list of [B,C,H,W] tensors
        batch_size = features[0].size(0)
        h, w = features[0].size()[2:]
        
        # 将特征reshape为N×HWC矩阵
        feat_matrix = torch.stack([
            f.view(batch_size, -1) for f in features
        ], dim=1)  # [B,N,HWC]
        
        # 计算层间相关性
        correlation = torch.bmm(
            feat_matrix.transpose(1,2), 
            feat_matrix
        )  # [B,HWC,N] x [B,N,HWC] -> [B,HWC,HWC]
        
        attention = torch.softmax(correlation, dim=-1)
        
        # 应用注意力权重
        attended_feats = torch.bmm(
            feat_matrix, 
            attention
        ).view(batch_size, self.n_resgroups, self.channels, h, w)
        
        # 加权求和
        output = torch.stack(features, dim=1) + self.alpha * attended_feats
        return output.mean(dim=1)  # [B,C,H,W]

3.2 通道空间注意力模块(CSAM)

CSAM模块创新性地使用3D卷积同时处理通道和空间维度：

python复制class CSAM(nn.Module):
    def __init__(self, channels=64):
        super().__init__()
        self.conv3d = nn.Conv3d(
            1, 1, kernel_size=(3,3,3),
            padding=(1,1,1)
        )
        self.sigmoid = nn.Sigmoid()
        self.beta = nn.Parameter(torch.zeros(1))
        
    def forward(self, x):
        # x: [B,C,H,W]
        b, c, h, w = x.size()
        
        # 添加维度并应用3D卷积
        x_3d = x.view(b, 1, c, h, w)
        attention = self.sigmoid(
            self.conv3d(x_3d)
        ).view(b, c, h, w)
        
        return x + self.beta * (x * attention)

4. 完整网络集成与训练策略

将各个模块组合成完整的HAN网络，并实现上采样和重建部分：

python复制class HAN(nn.Module):
    def __init__(self, scale_factor=4, n_resgroups=10, n_resblocks=20):
        super().__init__()
        self.sfe = ShallowFeatureExtraction()
        
        # 残差组
        self.resgroups = nn.ModuleList([
            ResidualGroup(n_resblocks) for _ in range(n_resgroups)
        ])
        
        # 注意力模块
        self.lam = LAM(n_resgroups)
        self.csam = CSAM()
        
        # 上采样
        self.upsample = nn.Sequential(
            nn.Conv2d(64, 64*scale_factor**2, 3, padding=1),
            nn.PixelShuffle(scale_factor),
            nn.Conv2d(64, 3, 3, padding=1)
        )
        
    def forward(self, x):
        # 浅层特征
        x = self.sfe(x)
        
        # 残差组处理
        features = []
        for group in self.resgroups:
            x = group(x)
            features.append(x)
        
        # 层注意力
        x = self.lam(features)
        
        # 通道空间注意力
        x = self.csam(x)
        
        # 上采样重建
        return self.upsample(x)

训练过程中有几个关键技巧值得注意：

• 学习率策略：采用余弦退火配合热启动
• 损失函数：组合使用L1损失和感知损失
• 数据增强：随机旋转、翻转和色彩抖动

python复制optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(
    optimizer, T_0=10, T_mult=2
)

criterion = nn.L1Loss()
perceptual_loss = PerceptualLoss()  # 预训练的VGG网络

for epoch in range(100):
    for lr, hr in dataloader:
        sr = model(lr)
        loss = criterion(sr, hr) + 0.1*perceptual_loss(sr, hr)
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    scheduler.step()

5. 评估与结果分析

在Set5、Set14等标准测试集上的评估是验证模型性能的关键。我们实现了与原始论文相同的评估流程：

python复制def evaluate(model, test_loader, device):
    model.eval()
    psnr_values = []
    ssim_values = []
    
    with torch.no_grad():
        for lr, hr in test_loader:
            lr, hr = lr.to(device), hr.to(device)
            sr = model(lr)
            
            # 计算PSNR和SSIM
            psnr = 10 * torch.log10(1 / torch.mean((sr - hr)**2))
            ssim = structural_similarity(
                sr.cpu().numpy(), 
                hr.cpu().numpy(),
                multichannel=True
            )
            
            psnr_values.append(psnr.item())
            ssim_values.append(ssim)
    
    return np.mean(psnr_values), np.mean(ssim_values)

在实现过程中，有几个容易出错的点需要特别注意：

1. LAM中的矩阵乘法顺序：确保特征矩阵转置正确
2. CSAM的3D卷积维度：输入必须是5D张量
3. 上采样比例因子：需要与数据预处理保持一致

经过完整训练后，在Set14数据集上（×4超分）的典型结果如下：

在实际项目中，我发现调整残差组的数量和CSAM的位置会对最终性能产生显著影响。特别是在资源有限的情况下，适当减少残差组数量（如从10个减到6个）可以在保持较好性能的同时大幅降低计算开销。另一个实用技巧是在训练后期逐步增加感知损失的权重，这有助于提升视觉质量。