ICCV 2021 | 最快视觉Transformer！Facebook提出LeViT：快速推理的视觉Transformer
点击下方卡片，关注“CVer”公众号
AI/CV重磅干货，第一时间送达
本文转载自：集智书童
LeViT: a Vision Transformer in ConvNet’s Clothing for Faster Inference
论文：https://arxiv.org/abs/2104.01136
代码（刚刚开源）：
https://github.com/facebookresearch/LeViT
吸取CNN优点！LeViT：快速推理的视觉Transformer,在速度/准确性的权衡方面LeViT明显优于现有的CNN和视觉Transformer，比如ViT、DeiT等，而且top-1精度为80％的情况下LeViT比CPU上的EfficientNet快3.3倍。
作者单位：Facebook
1 简介
本文的工作利用了基于注意力体系结构中的最新发现，该体系结构在高度并行处理硬件上具有竞争力。作者从卷积神经网络的大量文献中重新评估了原理，以将其应用于Transformer，尤其是分辨率降低的激活图。同时作者还介绍了Attention bias，一种将位置信息集成到视觉Transformer中的新方法。
图1 LeViT性能对比
最终作者提出了LeVIT：一种用于快速推理的混合神经网络。考虑在不同的硬件平台上采用不同的效率衡量标准，以最好地反映各种应用场景。作者通过广泛的实验表明该方法适用于大多数体系结构。总体而言，在速度/准确性的权衡方面，LeViT明显优于现有的卷积网络和视觉Transformer。例如，在ImageNet Top-1精度为80％的情况下，LeViT比CPU上的EfficientNet快3.3倍。
相同计算复杂度的情况下Transformer为什么快？
大多数硬件加速器(gpu，TPUs)被优化以用来执行大型矩阵乘法。在Transformer中，注意力机制和MLP块主要依靠这些操作。相比之下，卷积需要复杂的数据访问模式，因此它们的操作通常受io约束。这些考虑对于我们探索速度/精度的权衡是很重要的。
本文主要贡献：
采用注意力机制作为下采样机制的multi-stage transformer 结构;
一种计算效率高的patch descriptor，可以减少第一层特征的数量;
使用Translation-invariant attention bias取代ViT中的位置嵌入;
为了提高给定计算时间的网络容量，作者重新设计了Attention-MLP Block。
2 LeViT的设计
2.1 LeViT设计原则
LeViT以ViT的架构和DeiT的训练方法为基础，合并了对卷积架构有用的组件。第1步是获得Compatible Representation。如果不考虑classification embedding的作用，ViT就是一个处理激活映射的Layer的堆叠。
实际上，中间“Token”嵌入可以看作是FCN体系结构中传统的C×H×W激活映射(BCHW格式)。因此，适用于激活映射(池、卷积)的操作可以应用于DeiT的中间表征。
LeViT优化了计算体系结构，不一定是为了最小化参数的数量。ResNet系列比VGG更高效的设计原则之一是在其前2个阶段使用相对较小的计算预算应用strong resolution reductions。当激活映射到达ResNet的第3阶段时，其分辨率已经缩小到足以将卷积应用于小的激活映射，从而降低了计算成本。
2.2 LeViT组件
1、Patch embedding
初步分析表明，在transformer组的输入上应用一个小卷积可以提高精度。因此在LeViT中作者选择对输入应用4层3×3卷积(stride2)来降低分辨率。channel的数量是C=3,32,64,128,256。
以上操作减少了对transformer下层的激活映射的输入，同时不丢失重要信息。LeViT-256的patch extractor用184 MFLOPs将图像形状(3,224,224)转换为(256,14,14)。作为比较，ResNet-18的前10层使用1042 MFLOPs执行相同的dimensionality reduction。
为什么在transformer组的输入上应用一个小卷积可以提高精度？
2、No classification token
为了使用BCHW张量形式，LeViT删除了classification token。类似于卷积网络，在最后一个激活映射上使用GAP来代替，这将产生一个用于分类器的embedding。在训练中进行蒸馏，作者分别训练分类和蒸馏的Head。在测试时，平均2个Head的输出。在实践中，LeViT可以使用BNC或BCHW张量格式。
3、Normalization layers and activations
ViT架构中的FC层相当于1x1卷积。ViT在每个注意点和MLP单元之前使用层归一化。对于LeViT，每次卷积之后都要进行BN操作。然后与residual connection连接起来的每个BN权重参数初始化为零。BN可以与之前的卷积合并来进行推理，这比层归一化有运行优势(例如，在EfficientNet B0上，这种融合将GPU的推理速度提高了2倍)。而DeiT使用GELU函数，而LeViT的非线性激活都是Hardswish。
class Linear_BN(torch.nn.Sequential):
def __init__(self, a, b, bn_weight_init=1, resolution=-100000):
super().__init__()
self.add_module('c', torch.nn.Linear(a, b, bias=False))
bn = torch.nn.BatchNorm1d(b)
torch.nn.init.constant_(bn.weight, bn_weight_init)
torch.nn.init.constant_(bn.bias, 0)
self.add_module('bn', bn)
global FLOPS_COUNTER
output_points = resolution**2
FLOPS_COUNTER += a * b * output_points
@torch.no_grad()
def fuse(self):
l, bn = self._modules.values()
w = bn.weight / (bn.running_var + bn.eps)**0.5
w = l.weight * w[:, None]
b = bn.bias - bn.running_mean * bn.weight / \
(bn.running_var + bn.eps)**0.5
m = torch.nn.Linear(w.size(1), w.size(0))
m.weight.data.copy_(w)
m.bias.data.copy_(b)
return m
def forward(self, x):
l, bn = self._modules.values()
x = l(x)
return bn(x.flatten(0, 1)).reshape_as(x)
4、Multi-resolution pyramid
LeViT在transformer架构中集成了ResNet stage。在各个stage中，该体系结构类似于一个visual transformer:一个带有交替MLP和激活块的残差模块。下面是注意块的修改。
class Attention(torch.nn.Module):
def __init__(self, dim, key_dim, num_heads=8,
attn_ratio=4,
activation=None,
resolution=14):
super().__init__()
self.num_heads = num_heads
self.scale = key_dim ** -0.5
self.key_dim = key_dim
self.nh_kd = nh_kd = key_dim * num_heads
self.d = int(attn_ratio * key_dim)
self.dh = int(attn_ratio * key_dim) * num_heads
self.attn_ratio = attn_ratio
h = self.dh + nh_kd * 2
self.qkv = Linear_BN(dim, h, resolution=resolution)
self.proj = torch.nn.Sequential(activation(), Linear_BN(
self.dh, dim, bn_weight_init=0, resolution=resolution))
points = list(itertools.product(range(resolution), range(resolution)))
N = len(points)
attention_offsets = {}
idxs = []
for p1 in points:
for p2 in points:
offset = (abs(p1[0] - p2[0]), abs(p1[1] - p2[1]))
if offset not in attention_offsets:
attention_offsets[offset] = len(attention_offsets)
idxs.append(attention_offsets[offset])
self.attention_biases = torch.nn.Parameter(
torch.zeros(num_heads, len(attention_offsets)))
self.register_buffer('attention_bias_idxs',
torch.LongTensor(idxs).view(N, N))
global FLOPS_COUNTER
#queries * keys
FLOPS_COUNTER += num_heads * (resolution**4) * key_dim
# softmax
FLOPS_COUNTER += num_heads * (resolution**4)
#attention * v
FLOPS_COUNTER += num_heads * self.d * (resolution**4)
@torch.no_grad()
def train(self, mode=True):
super().train(mode)
if mode and hasattr(self, 'ab'):
del self.ab
else:
self.ab = self.attention_biases[:, self.attention_bias_idxs]
def forward(self, x): # x (B,N,C)
B, N, C = x.shape
qkv = self.qkv(x)
q, k, v = qkv.view(B, N, self.num_heads, -
1).split([self.key_dim, self.key_dim, self.d], dim=3)
q = q.permute(0, 2, 1, 3)
k = k.permute(0, 2, 1, 3)
v = v.permute(0, 2, 1, 3)
attn = (
(q @ k.transpose(-2, -1)) * self.scale
+
(self.attention_biases[:, self.attention_bias_idxs]
if self.training else self.ab)
)
attn = attn.softmax(dim=-1)
x = (attn @ v).transpose(1, 2).reshape(B, N, self.dh)
x = self.proj(x)
return x
5、Downsampling
在LeViT stage之间，一个缩小的注意块减少了激活映射的大小:在Q转换之前应用一个subsampling，然后传播到soft activation的输出。这将一个大小为 的输入张量映射到一个大小为 的输出张量。由于尺度的变化这个注意块的使用没有残差连接。同时为了防止信息丢失，这里将注意力头的数量设为 。
class Subsample(torch.nn.Module):
def __init__(self, stride, resolution):
super().__init__()
self.stride = stride
self.resolution = resolution
def forward(self, x):
B, N, C = x.shape
x = x.view(B, self.resolution, self.resolution, C)[
:, ::self.stride, ::self.stride].reshape(B, -1, C)
return x
6、Attention bias instead of a positional embedding
在transformer架构中的位置嵌入是一个位置依赖可训练的向量，在将token嵌入输入到transformer块之前，将其添加到token嵌入。如果没有它，转换器输出将独立于输入标记的排列。位置嵌入的Ablations会导致分类精度的急剧下降。
然而，位置嵌入只包含在注意块序列的输入上。因此，由于位置编码对higher layer也很重要，所以它很可能仍然处于中间表示中。
因此，LeViT在每个注意块中提供位置信息，并在注意机制中明确地注入相对位置信息:只是在注意力图中添加了注意偏向。对于每个head ，每2个像素 和 之间的标量值计算方式为：
第一项是经典的注意力。第二个是translation-invariant attention bias。每个Head有H×W参数对应不同的像素偏移量。对称差异 和 鼓励用 flip invariance进行训练。
self.attention_biases = torch.nn.Parameter(
torch.zeros(num_heads, len(attention_offsets)))
7、Smaller keys
由于translation-invariant attention bias偏置项减少了key对位置信息编码的压力，因此LeViT减少了key矩阵相对于V矩阵的大小。如果key大小为 , V则有2D通道。key的大小可以减少计算key product 所需的时间。
对于没有残差连接的下采样层，将V的维数设置为4D，以防止信息丢失。
8、Attention activation
在使用常规线性投影组合不同Heads的输出之前，对product 应用Hardswish激活。这类似于ResNet bottleneck residual block，V是一个1×1卷积的输出， 对应一个spatial卷积，projection是另一个1×1卷积。
9、Reducing the MLP blocks
在ViT中，MLP residual块是一个线性层，它将嵌入维数增加了4倍，然后用一个非线性将其减小到原来的嵌入维数。但是对于视觉架构，MLP通常在运行时间和参数方面比注意Block更昂贵。
对于LeViT, MLP是1x1卷积，然后是通常的BN。为了减少计算开销，将卷积的展开因子从4降低到2。一个设计目标是注意力和MLP块消耗大约相同数量的FLOPs。
2.3 LeViT家族
3 实验
3.1 速度对比
ResNet50的精度，但是是起飞的速度。
3.2 SOTA对比
论文PDF和代码下载
后台回复：LeViT，即可下载上述论文和代码
后台回复：CVPR2021，即可下载CVPR 2021论文和代码开源的论文合集
后台回复：Transformer综述，即可下载最新的两篇Transformer综述PDF
重磅！Transformer交流群成立
扫码添加CVer助手，可申请加入CVer-Transformer 微信交流群，方向已涵盖：目标检测、图像分割、目标跟踪、人脸检测&识别、OCR、姿态估计、超分辨率、SLAM、医疗影像、Re-ID、GAN、NAS、深度估计、自动驾驶、强化学习、车道线检测、模型剪枝&压缩、去噪、去雾、去雨、风格迁移、遥感图像、行为识别、视频理解、图像融合、图像检索、论文投稿&交流、PyTorch和TensorFlow等群。
一定要备注：研究方向+地点+学校/公司+昵称（如Transformer+上海+上交+卡卡），根据格式备注，可更快被通过且邀请进群
▲长按加小助手微信，进交流群
▲点击上方卡片，关注CVer公众号
整理不易，请点赞和在看