自注意力机制在图像中的应用

笔者的毕设是做人脸图像的补全，开始使用经典的变分自编码器模型，能达到比较好的补全效果．后来看到BIGGAN的论文，里边他们使用了self attention提高图片生成的效果，查阅了相关资料后也在模型中加入了自注意力层，确实对补全后的图像有了显著的提升．当然了BIGGAN生成的图片能达到以假乱真的地步也不光是靠加入了自注意力层，具体可以看http://www.twistedwg.com/2018/10/02/BigGAN.html这篇博客里对BIGGAN的解读．本文主要分析为何加入自注意力层可以提高生成图片的质量以及如何实现自注意力层．

BIGGAN将自注意力层加入到模型中主要是参考的Self Attention Gan(SAGAN)，它将self attention layer 加入到GAN中，使得编码器和解码器能够更好地对图像不同区域的关系进行建模．

经典的DCGAN(Deep Convolutional GAN)采用多层卷积网络构建编码器和解码器，但是这种网络的表征能力受限于邻域的 kernel size，就算后期感受野越来越大，终究还是局部区域的运算，这样就忽略了全局其他片区（比如很远的像素）对当前区域的贡献。

在图像领域，自注意力机制会学习某一像素点和其他所有位置（包括较远位置）的像素点之间的关系，即捕获long-range的关系．也就是说对于SAGAN，它会利用所有位置的特征来帮助生成图片的某一细节，使生成的图片更加逼真．

　　　　　　　　　　　　　　　　　　　　　　　　

 

对于ＳＡＧＡＮ，首先对输入的 feature map x 进行线性映射得到  f ,  g,  h, 其中 f(x) = W_f *x,  g(x) = W_g*x. 这两个特征空间中的值将用来计算attention　β_j,i  = softmax( f(x_i)^T g(x_i))

                                                                                                                                                             

 

β_j,i 用来表示i^th位置对生成j^th 区域的关系权重，Ｎ是特征位置的数目．将　h　与attention点乘得到注意力层的输出结果o = (o₁, o_2, ..., o_{j ,}...o_n)

                                                                                                                                            

　

　　　　　　　　　　　　　　　　　　

　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　SAGAN中的自注意力机制

                                                                                                                                           

 

其中W_f Wg W_h 都是1×1的卷积，　最后将attention layer的结果乘以一个系数再加上原来的特征图就是最后的结果 y_i = x_i + γo_i. 其中γ是一个可学习的参数，在训练过程中初始化为０，网络在训练初始阶段主要依赖邻域特征，之后慢慢的增大对对较远区域依赖的权重．

 

BIGGAN实现的SelfAttention代码如下

def hw_flatten(x) :
    return tf.reshape(x, shape=[x.shape[0], -1, x.shape[-1]])

def self_attention(x, channels, sn=False, scope='self_attention'):
    with tf.variable_scope(scope):
        f = conv(x, channels // 8, kernel=1, stride=1, scope='f_conv')  # [bs, h, w, c']
        g = conv(x, channels // 8, kernel=1, stride=1, scope='g_conv')  # [bs, h, w, c']
        h = conv(x, channels, kernel=1, stride=1, sn=sn, scope='h_conv')  # [bs, h, w, c]

        # N = h * w
        s = tf.matmul(hw_flatten(g), hw_flatten(f), transpose_b=True)  # # [bs, N, N]

        beta = tf.nn.softmax(s)  # attention map

        o = tf.matmul(beta, hw_flatten(h))  # [bs, N, C]
        gamma = tf.get_variable("gamma", [1], initializer=tf.constant_initializer(0.0))

        o = tf.reshape(o, shape=x.shape)  # [bs, h, w, C]
        x = gamma * o + x

    return x

 

reference -  "Self-Attention Generative Adversarial Networks"  Han Zhang,  Ian Goodfellow,  Dimitris Metaxas,  Augustus Odena, 2019

                   "Attention？ Attention!"  Lilian Weng 

                   "综述---图像处理中的注意力机制"  https://blog.csdn.net/xys430381_1/article/details/89323444