参考目录:
首先,这个论文中的模块,叫做non-local block,然后这个思想是基于NLP中的self-attention自注意力机制的。所以在提到CV中的self-attention,最先想到的就是non-local这个论文。这个论文提出的动机如下:
卷积运算和递归操作都在空间或时间上处理一个local邻域;只有在重复应用这些运算、通过数据逐步传播信号时,才能捕获long-range相关性。
换句话说,在卷积网络中,想要增加视野域,就要不断的增加卷积层数量和池化层数量,换句话说,增加视野域就是增加网络的深度。这样必然会增加计算的成本,参数的数量,还需要考虑梯度消失问题。
本次我们学习先看代码,然后再从论文中解析代码。
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(self,x):
"""
inputs :
x : input feature maps( B X C X W X H)
returns :
out : self attention value + input feature
attention: B X N X N (N is Width*Height)
"""
m_batchsize,C,width ,height = x.size()
proj_query = self.query_conv(x).view(m_batchsize,-1,width*height).permute(0,2,1) # B X CX(N)
proj_key = self.key_conv(x).view(m_batchsize,-1,width*height) # B X C x (*W*H)
energy = torch.bmm(proj_query,proj_key) # transpose check
attention = self.softmax(energy) # BX (N) X (N)
proj_value = self.value_conv(x).view(m_batchsize,-1,width*height) # B X C X N
out = torch.bmm(proj_value,attention.permute(0,2,1) )
out = out.view(m_batchsize,C,width,height)
out = self.gamma*out + x
return out,attention
输入特征图为BatchxChannelxHeightxWidth,我们先把这个输入特征图x分别放入:
我们要逐个像素的计算query和key的相似度,然后相似度高的像素更为重要,相似度低的像素就不那么重要,每个像素我们用channel//8这个长度的向量来表示。(这里可能比较抽象,毕竟self-attention的原版是NLP领域的,non-local是从NLP中照搬过来的,所以不太好直接理解)
相似度计算是通过向量的乘法来表示的,那么我们肯定不能把这个HeightxWidth这么多像素一个一个计算像素的相似度。所以我们把BatchxChannel//8xHeightxWidth转换成BatchxChannel//8xN的形式,这里的N是HeightxWidth,N表示图中像素的数量。
然后我们用torch.bmm()
来做矩阵的乘法:(N,Channel//8)和(Channel//8,N)两个矩阵相乘,得到一个(N,N)的矩阵。
这个(N,N)矩阵中的第i行第j列元素的值,是图中i位置像素和j位置像素的相关性!然后我们把value矩阵和这个(N,N)再进行一次矩阵乘法,这样得到的输出,就是考虑了全局信息的特征图了。
第二次矩阵乘法中,是(Channel,N)和(N,N)的相乘,得到的输出的特征图中的每一个值,都是N个值的加权平均,这也说明了输出的特征图中的每一个值,都是考虑了整张图的像素的。
上图是论文中对于non-local的结构图。可以看到,先通过1x1的卷积,降低一下通道数,然后通过
分别是query和key,然后这两个卷积得到(N,N)的矩阵,然后再与
(value)进行矩阵乘法。
好吧我承认和代码在通道数上略微有些出入,但是大体思想相同。