SA-GAN： self-attention 的 pytorch 实现（针对图像）

最新推荐文章于 2025-07-09 15:12:08 发布

月下花弄影

最新推荐文章于 2025-07-09 15:12:08 发布

阅读量7k

点赞数 1

CC 4.0 BY-SA版权

分类专栏：数学理论 GAN pytorch

本文链接：https://blog.youkuaiyun.com/qq_27261889/article/details/87626592

数学理论同时被 3 个专栏收录

15 篇文章

订阅专栏

pytorch

15 篇文章

订阅专栏

GAN

8 篇文章

订阅专栏

本文探讨了基于条件的卷积GAN在生成特定类型图像时遇到的挑战，尤其是在处理复杂全局结构和细密纹理方面。提出了自我注意力GAN（SA-GAN）作为解决方案，通过引入全局依赖性来改善生成质量，确保不同位置特征的有效结合。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

问题

基于条件的卷积GAN 在那些约束较少的类别中生成的图片较好，比如大海，天空等；但是在那些细密纹理，全局结构较强的类别中生成的图片不是很好，如人脸（可能五官不对应），狗（可能狗腿数量有差，或者毛色不协调）。

可能的原因

大部分卷积神经网络都严重依赖于局部感受野，而无法捕捉全局特征。另外，在多次卷积之后，细密的纹理特征逐渐消失。

SA-GAN解决思路

不仅仅依赖于局部特征，也利用全局特征，通过将不同位置的特征图结合起来（转置就可以结合不同位置的特征）。

"""
    Author: Mingle Xu
    Time: 18 Mar, 2019
    Time: 2020-04-08 22:04:07 revise something
"""

import torch
import torch.nn as nn


class self_attention(nn.Module):
    r"""
        Create global dependence.
        Source paper: https://arxiv.org/abs/1805.08318
    """

    def __init__(self, in_channles):
        super(self_attention, self).__init__()
        self.in_channels = in_channles

        self.f = nn.Conv2d(in_channels=in_channles, out_channels=in_channles // 8, kernel_size=1)
        self.g = nn.Conv2d(in_channels=in_channles, out_channels=in_channles // 8, kernel_size=1)
        self.h = nn.Conv2d(in_channels=in_channles, out_channels=in_channles, kernel_size=1)
        self.softmax_ = nn.Softmax(dim=2)
        self.gamma = nn.Parameter(torch.zeros(1))

        self.init_weight(self.f)
        self.init_weight(self.g)
        self.init_weight(self.h)

    def forward(self, x):
        batch_size, channels, height, width = x.size()

        assert channels == self.in_channels

        f = self.f(x).view(batch_size, -1, height * width).permute(0, 2, 1)  # B * (H * W) * C//8
        g = self.g(x).view(batch_size, -1, height * width)  # B * C//8 * (H * W)

        attention = torch.bmm(f, g)  # B * (H * W) * (H * W)
        attention = self.softmax_(attention)

        h = self.h(x).view(batch_size, channels, -1)  # B * C * (H * W)

        self_attention_map = torch.bmm(h, attention).view(batch_size, channels, height, width)  # B * C * H * W

        return self.gamma * self_attention_map + x

    def init_weight(self, conv):
        nn.init.kaiming_uniform_(conv.weight)
        if conv.bias is not None:
            conv.bias.data.zero_()