Self-Attention GAN 中的 self-attention 机制

作者丨尹相楠

学校丨里昂中央理工博士在读

研究方向丨人脸识别、对抗生成网络

Self Attention GAN 用到了很多新的技术。最大的亮点当然是 self-attention 机制，该机制是 Non-local Neural Networks [1] 这篇文章提出的。其作用是能够更好地学习到全局特征之间的依赖关系。因为传统的 GAN 模型很容易学习到纹理特征：如皮毛，天空，草地等，不容易学习到特定的结构和几何特征，例如狗有四条腿，既不能多也不能少。 

除此之外，文章还用到了 Spectral Normalization for GANs [2] 提出的谱归一化。谱归一化的解释见本人这篇文章：详解GAN的谱归一化（Spectral Normalization）。

但是，该文代码中的谱归一化和原始的谱归一化运用方式略有差别： 

1. 原始的谱归一化基于 W-GAN 的理论，只用在 Discriminator 中，用以约束 Discriminator 函数为 1-Lipschitz 连续。而在 Self-Attention GAN 中，Spectral Normalization 同时出现在了 Discriminator 和 Generator 中，用于使梯度更稳定。除了生成器和判别器的最后一层外，每个卷积/反卷积单元都会上一个 SpectralNorm。 

2. 当把谱归一化用在 Generator 上时，同时还保留了 BatchNorm。Discriminator 上则没有 BatchNorm，只有 SpectralNorm。 

3. 谱归一化用在 Discriminator 上时最后一层不加 Spectral Norm。 

最后，self-attention GAN 还用到了 cGANs With Projection Discriminator 提出的 conditional normalization 和 projection in the discriminator。这两个技术我还没有来得及看，而且 PyTorch 版本的 self-attention GAN 代码中也没有实现，就先不管它们了。

本文主要说的是 self-attention 这部分内容。

▲ 图1. Self-Attention

Self-Attention

在卷积神经网络中，每个卷积核的尺寸都是很有限的（基本上不会大于 5），因此每次卷积操作只能覆盖像素点周围很小一块邻域。

对于距离较远的特征，例如狗有四条腿这类特征，就不容易捕获到了（也不是完全捕获不到，因为多层的卷积、池化操作会把 feature map 的高和宽变得越来越小，越靠后的层，其卷积核覆盖的区域映射回原图对应的面积越大。但总而言之，毕竟还得需要经过多层映射，不够直接）。

Self-Attention 通过直接计算图像中任意两个像素点之间的关系，一步到位地获取图像的全局几何特征。 

论文中的公式不够直观，我们直接看文章的 PyTorch 的代码，核心部分为 sagan_models.py：

class Self_Attn(nn.Module):
    """ Self attention Layer"""
    def __init__(self,in_dim,activation):
        super(Self_Attn,self).__init__()
        self.chanel_in = in_dim
        self.activation = activation

        self.query_conv = nn.Conv2d(in_channels = in_dim , out_channels = in_dim//8 , kernel_size= 1)
        self.key_conv = nn.Conv2d(in_channels = in_dim , out_channels = in_dim//8 , kernel_size= 1)
        self.value_conv = nn.Conv2d(in_channels = in_dim , out_channels = in_dim , kernel_size= 1)
        self.gamma = nn.Parameter(torch.zeros(1))

        self.softmax  = nn.Softmax(dim=-1) #
    def forward