Non_Local 网络模块的解析_non-local模块

该模块可以插入到现有的CNN网络结构中，优点是可以提高整体的性能，缺点是参数多，计算量大

写此文章只为自己日后回顾

首先base cnn 为resnet50 ,插入在block3之后，模块的输入和输出都不变，这也是可以即插即用的原因。

首先输入为【B,H,W,C】 先对输入使用1X1大小的卷积核做卷积，降低输入通道。减少计算量。然后把【B,H,W,C】reshape 成[B,HW,C]，然后两个相乘（其中一个transpose）,这样可以得到【B,HW,HW】,可以得到图像像素和其他位置的相关性，然后将结果做softmax 处理，突出共性，然后将softmax得到的记过和【B,HW,C】矩阵相乘，将权重应用到输入上，然后在和原始输入相加，就完成位置注意力机制。

这个模块，可以让模型把注意力放在要识别的物体上。详细代码如下：

def nonlocal_dot(net, depth, embed=True, softmax=False, maxpool=2, scope=None):
  """ Implementation of the non-local block in its various forms.
  See "Non-local Neural Networks" by
  Xiaolong Wang, Ross Girshick, Abhinav Gupta, Kaiming He
  https://arxiv.org/pdf/1711.07971.pdf
  Args:
  - `net`: The symbolic input into the block, a (B,H,W,C) Tensor.
  - `depth`: The number of channels in which to execute the non-local operation.
  - `embed`: Whether or not use the "embedded version" as in Sec.3.2
  - `softmax`: Whether or not to use the softmax operation which makes it
               equivalent to soft-attention.
  - `maxpool`: How large of a max-pooling (Sec.3.3) to use to help reduce
               the computational burden. Default is 2, use `False` for none.
  - `scope`: An optional scope for all created variables.
  Returns:
    The symbolic output of the non-local block operation.
  Note:
    The final BatchNorm's gamma is initialized to zero, so as to make this a
    no-op (skip) at initialization, as described in Sec.4.1.
  """
  with tf.variable_scope(scope, 'nonlocal', values=[net]) as sc:
    with slim.arg_scope([slim.conv2d], normalizer_fn=None):
      if embed:
        #change input channels to 512
        a = conv2d_same(net, depth, 1, stride=1, scope='embA')
        b = conv2d_same(net, depth, 1, stride=1, scope='embB')
      else:
        a, b = net, net
 
      g_orig = g = conv2d_same(net, depth, 1, stride=1, scope='g')
 
    if maxpool is not False and maxpool > 1:
      b = slim.max_pool2d(b, [maxpool, maxpool], stride=maxpool, scope='pool')
      g = slim.max_pool2d(g, [maxpool, maxpool], stride=maxpool, scope='pool')
 
    # Flatten from (B,H,W,C) to (B,HW,C) or similar
    a_flat = tf.reshape(a, [tf.shape(a)[0], -1, tf.shape(a)[-1]])
    b_flat = tf.reshape(b, [tf.shape(b)[0], -1, tf.shape(b)[-1]])
    g_flat = tf.reshape(g, [tf.shape(g)[0], -1, tf.shape(g)[-1]])
 
    a_flat.set_shape([8, a.shape[1] * a.shape[2] if None not in a.shape[1:3] else None, a.shape[-1]])
    b_flat.set_shape([8, b.shape[1] * b.shape[2] if None not in b.shape[1:3] else None, b.shape[-1]])
    g_flat.set_shape([8, g.shape[1] * g.shape[2] if None not in g.shape[1:3] else None, g.shape[-1]])
 
    # Compute f(a, b) -> (B,HW,HW) 计算相似性
 
    a_flat_new = tf.gather(a_flat,0)
    b_flat_new = tf.gather(b_flat,0)
    
    print("&&&&&&&&&&&&&&&&&&&&&&&&&")
    print(a_flat_new.shape) #[49,512]
    print(b_flat_new.shape) #[16,512]
    f0 = tf.matmul(a_flat_new, tf.transpose(b_flat_new, [1,0]))
    f0 = tf.reshape(f0,[-1,f0.shape[0],f0.shape[1]])
    print("f0.shape") #[1,49,16]
    print(f0.shape) 
  
    a_flat_new = tf.gather(a_flat,1)
    b_flat_new = tf.gather(b_flat,1)
    print("&&&&&&&&&&&&&&&&&&&&&&&&&")
    print(a_flat_new.shape)
    print(b_flat_new.shape)
    f1 = tf.matmul(a_flat_new, tf.transpose(b_flat_new, [1,0]))
    f1 = tf.reshape(f1,[-1,f1.shape[0],f1.shape[1]])
    print("f1.shape")
    print(f1.shape)
    
    f = tf.concat([f0,f1],axis=0)
    print("f.shape")
    print(f.shape)
 
    for i in range(6):
      i = i+2
      a_flat_new = tf.gather(a_flat,i)
      b_flat_new = tf.gather(b_flat,i)
      print("&&&&&&&&&&&&&&&&&&&&&&&&&")
      print(a_flat_new.shape)
      print(b_flat_new.shape)
      f0 = tf.matmul(a_flat_new, tf.transpose(b_flat_new, [1,0]))
      f0 = tf.reshape(f0,[-1,f0.shape[0],f0.shape[1]])
      print("f0.shape")
      print(f0.shape)
      f =tf.concat([f,f0],axis=0)
      
    print("f.shape")
    print(f.shape) #[8,49,16]
 
 
    if softmax:
        f = tf.nn.softmax(f)
    else:
        f = f / tf.cast(tf.shape(f)[-1], tf.float32)
    # Compute f * g ("self-attention") -> (B,HW,C)
    
    print("********************")
    print(g_flat.shape) #[8,16,512]
    
    f_flat_new = tf.gather(f,0)
    g_flat_new = tf.gather(g_flat,0)
    print("###################")
    print(f_flat_new.shape) #[49,16]
    print(g_flat_new.shape) #[16,512]
    f0 = tf.matmul(f_flat_new, g_flat_new)
    f0 = tf.reshape(f0,[-1,f0.shape[0],f0.shape[1]])
    print("f0.shape") #[1,49,16]
    print(f0.shape) 
  
    f_flat_new = tf.gather(f,1)
    g_flat_new = tf.gather(g_flat,1)
    print("##########################")
    print(f_flat_new.shape)
    print(g_flat_new.shape)
    f1 = tf.matmul(f_flat_new, g_flat_new)
    f1 = tf.reshape(f1,[-1,f1.shape[0],f1.shape[1]])
    print("f1.shape")
    print(f1.shape)
    
    f_new = tf.concat([f0,f1],axis=0)
    print("f_new.shape")
    print(f_new.shape)
 
    for i in range(6):
      i = i+2
      f_flat_new = tf.gather(f,i)
      g_flat_new = tf.gather(g_flat,i)
      print("&&&&&&&&&&&&&&&&&&&&&&&&&")
      print(f_flat_new.shape)
      print(g_flat_new.shape)
      f0 = tf.matmul(f_flat_new, g_flat_new)
      f0 = tf.reshape(f0,[-1,f0.shape[0],f0.shape[1]])
      print("f0.shape")
      print(f0.shape)
      f_new =tf.concat([f_new,f0],axis=0)
 
    print("f_new.shape")
    print(f_new.shape) #[8,49,16]
    #fg = tf.matmul(f, g_flat)
    
    # Expand and fix the static shapes TF lost track of.
    fg = tf.reshape(f_new, tf.shape(g_orig))
    # fg.set_shape(g.shape)  # NOTE: This actually appears unnecessary.
 
    # Go back up to the original depth, add residually, zero-init.
    #with slim.arg_scope([slim.conv2d],
    #                    weights_initializer=tf.zeros_initializer()):
    with slim.arg_scope([slim.batch_norm], param_initializers={'gamma': tf.zeros_initializer()}):
      fg = conv2d_same(fg, net.shape[-1], 1, stride=1, scope='fgup')
    net = net + fg
 
    return slim.utils.collect_named_outputs(None, sc.name, net)