Tensorflow入门教程（三十二）——Non-Local VNet3D

近年来，自注意力机制已经广泛应用于自然语言处理中，今天我将分享一下如何将自注意力机制应用在图像分割任务中。

1、Attention背景介绍

Attention机制思想最早是在九几年就提出来了，但是真正火起来应该算是2014年google mind团队的这篇论文《Recurrent Models of Visual Attention》，他们在RNN模型上使用了attention机制来进行图像分类。随后，Bahdanau等人在论文《Neural Machine Translation by Jointly Learning to Align and Translate》中，使用类似attention的机制在机器翻译任务上将翻译和对齐同时进行，他们的工作算是第一个将attention机制应用到NLP领域中。接着attention机制被广泛应用在基于RNN/CNN等神经网络模型的各种NLP任务中。2017年，google机器翻译团队发表的《Attention is all you need》中大量使用了自注意力（self-attention）机制来学习文本表示。自注意力机制也成为近期的研究热点，并在各种NLP任务上进行探索。

2、Non-local Neural Networks

2018年由王小龙，何凯明发表的《Non-local Neural Networks》论文是将自注意力机制应用在计算机视觉领域的开山之作。详细内容请阅读原文。这里就简单说一下论文的核心思想：在计算机视觉中卷积操作是在局部区域进行特征提取操作的，然后再通过不断堆叠多个卷积层和下采样层（池化层或卷积层）模块来形成在更大的感受野上进行特征提取操作，但这存在3个不足：

(1) 捕获长范围依赖的效率太低；

(2) 由于网络很深，需要小心的设计模块和梯度；

(3) 当需要在比较远位置之间来回传递消息时，局部操作是很难实现的。

故作者提出非局部操作算子，可以捕获时间(一维时序信号)、空间(图片)和时空(视频序列)的长范围依赖。这样设计的好处是：

(1) 相比较于不断堆叠卷积，非局部操作直接计算两个位置(可以是时间位置、空间位置和时空位置)之间的关系即可快速捕获长范围依赖，但是会忽略其欧式距离。

(2) 非局部操作计算效率很高，要达到同等效果，只需要更少的堆叠层。

(3) 非局部操作可以保证输入尺度和输出尺度不变，这种设计可以很容易嵌入到目前的网络架构中。

论文中的非局部模块结构如上图所示。首先网络输入X= (T, H, W, 1024) ，经过权重矩阵变换

，

，论文中使用的是（1x1x1）大小的卷积核来进行卷积操作变换得到(T, H, W, 512), (T, H, W, 512), 目的是降低通道数量以减少计算量；然后对两路输出进行reshape操作变成(THW, 512)，对两路输出进行矩阵乘操作(其中一个要转置)，计算矩阵相似性，得到(THW，THW),再进行softmax操作，起到自注意力的作用；然后经过权重变换矩阵Wg，采用（1x1x1）的卷积操作，先通道降维，再reshape操作，与softmax输出结果进行矩阵乘操作，得到(THW, 512), 这一步是将自注意力机制应用到了所有通道的每张特征图对应位置上，本质上是输出的每个位置值都是由其他所有位置的加权平均；再经过一个1x1x1卷积恢复输出通道，保证输入输出尺度完全相同，最后输入与该输出进行残差连接输出。

3、Non-local VNet3D

VNet3D网络模型如下所示，论文中也提到Non-local模块是可以在现有的网络中即插即用的，但是Non-local模块中涉及到矩阵乘法计算，如果在VNet3D的浅层处添加Non-local模块，计算量非常大，硬件无法运行，论文中也建议将Non-local模块添加在深层处，例如可以在VNet3D网络中解码网络中的128通道和256通道所在层数添加Non-local模块来替换原来的卷积层。