CBAM生成对抗网络：从原理到实战的深度解析

1次阅读

没有评论

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

生成对抗网络（GAN）自 2014 年提出以来，在图像生成领域取得了巨大成功。然而，传统 GAN 在实际应用中存在几个明显问题：

训练不稳定 ：生成器和判别器的对抗训练容易陷入振荡，难以收敛
模式崩溃 ：生成器倾向于生成有限的几种样本，缺乏多样性
细节丢失 ：对图像局部特征捕捉能力不足，生成结果往往缺乏精细结构

这些问题很大程度上源于传统卷积操作的局限性——它平等对待所有空间位置和通道，无法自适应地关注重要区域。而注意力机制的引入，正是为了解决这一核心矛盾。

常见的注意力机制主要有三种：

SE（Squeeze-and-Excitation）模块 ：
仅考虑通道注意力
计算量小但空间信息利用不足
Non-local 网络 ：
捕获长距离依赖关系
计算复杂度高（O(N²)）
CBAM（Convolutional Block Attention Module）：
同时包含通道和空间注意力
计算效率高（O(C) + O(HW)）
可轻松嵌入现有网络

在 GAN 框架中，CBAM 尤其适合用于：
– 生成器的上采样层，增强细节生成能力
– 判别器的特征提取层，提升真假判别精度

CBAM 由两个顺序子模块组成：

通道注意力模块 ：
通过全局平均池化和最大池化获得通道统计信息
使用共享 MLP 生成通道权重
空间注意力模块 ：
沿通道维度应用平均池化和最大池化
通过卷积层生成空间权重图

import torch
import torch.nn as nn
import torch.nn.functional as F

class ChannelAttention(nn.Module):
    def __init__(self, in_planes, ratio=8):
        super(ChannelAttention, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.max_pool = nn.AdaptiveMaxPool2d(1)

        self.shared_MLP = nn.Sequential(nn.Conv2d(in_planes, in_planes // ratio, 1, bias=False),
            nn.ReLU(),
            nn.Conv2d(in_planes // ratio, in_planes, 1, bias=False)
        )
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        avg_out = self.shared_MLP(self.avg_pool(x))
        max_out = self.shared_MLP(self.max_pool(x))
        return self.sigmoid(avg_out + max_out)

class SpatialAttention(nn.Module):
    def __init__(self, kernel_size=7):
        super(SpatialAttention, self).__init__()
        assert kernel_size in (3,7), "kernel size must be 3 or 7"
        padding = 3 if kernel_size == 7 else 1

        self.conv = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        avg_out = torch.mean(x, dim=1, keepdim=True)
        max_out, _ = torch.max(x, dim=1, keepdim=True)
        x = torch.cat([avg_out, max_out], dim=1)
        x = self.conv(x)
        return self.sigmoid(x)

class CBAM(nn.Module):
    def __init__(self, planes):
        super(CBAM, self).__init__()
        self.ca = ChannelAttention(planes)
        self.sa = SpatialAttention()

    def forward(self, x):
        x = self.ca(x) * x
        x = self.sa(x) * x
        return x

# 数据预处理
transform = transforms.Compose([transforms.Resize(64),
    transforms.CenterCrop(64),
    transforms.ToTensor(),
    transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5))
])

# 初始化网络
generator = Generator(z_dim=100, img_channels=3)
discriminator = Discriminator(img_channels=3)
cbam = CBAM(64)  # 嵌入到生成器和判别器中

# 训练循环
for epoch in range(epochs):
    for i, real_imgs in enumerate(dataloader):
        # 训练判别器
        optimizer_D.zero_grad()
        z = torch.randn(batch_size, z_dim, 1, 1)
        fake_imgs = generator(z)

        real_validity = discriminator(real_imgs)
        fake_validity = discriminator(fake_imgs.detach())

        d_loss = -torch.mean(real_validity) + torch.mean(fake_validity)
        d_loss.backward()
        optimizer_D.step()

        # 训练生成器
        if i % 5 == 0:
            optimizer_G.zero_grad()
            gen_validity = discriminator(fake_imgs)
            g_loss = -torch.mean(gen_validity)
            g_loss.backward()
            optimizer_G.step()

在 CelebA 数据集上的对比实验显示：