【GiantPandaCV导语】基于Transformer的骨干网络，同时使用卷积与自注意力机制来保持全局性和局部性。模型在ResNet最后三个BottleNeck中使用了MHSA替换3x3卷积。属于早期的结合CNN+Transformer的工作。简单来讲Non-Local+Self Attention+BottleNeck = BoTNet
引言
本文的发展脉络如下图所示：
实际上沿着Transformer Block改进的方向进行的，与CNN架构也是兼容的。具体结构如下图所示：
两者都遵循了BottleNeck的设计原则，可以有效降低计算量。
设计Transformer中self attention存在几个挑战：
- 图片尺寸比较大，比如目标检测中分辨率在1024x1024
- 内存和计算量的占用高，导致训练开销比较大。
本文设计如下：
- 使用卷积识别底层特征的抽象信息。
- 使用self attention处理通过卷积层得到的高层信息。
这样可以有效处理大分辨率图像。
方法
BoTNet中MHSA模块如下图所示：
上边的这个MHSA Block是核心创新点，其与Transformer中的MHSA有所不同：
- 由于处理对象不是一维的，而是类似CNN模型，所以有非常多特性与此相关。
- 归一化这里并没有使用Layer Norm而是采用的Batch Norm，与CNN一致。
- 非线性激活，BoTNet使用了三个非线性激活
- 左侧content-position模块引入了二维的位置编码，这是与Transformer中最大区别。
由于该模块是处理BxCHW的形式，所以难免让人想起来Non Local 操作，这里列出笔者以前绘制的一幅图：
可以看出主要区别就是在于Content-postion模块引入的位置信息。
BoTNet细节设计：
整体的设计和ResNet50几乎一样，唯一不同在于最后一个阶段中三个BottleNeck使用了MHSA模块。具体这样做的原因是Self attention需要消耗巨大的计算量，在模型最后加入时候feature map的size比较小，相对而言计算量比较小。
实验
在目标检测和分割领域性能对比
分辨率改变对BoTNet帮助更大
消融实验-相对位置编码
BoTNet对ResNet系列模型的提升
BoTNet与更大的图片适配
BoTNet与Non-Local Net的比较
与ImageNet上结果比较
模型放缩的影响
显卡香气飘来，又是谷歌的骚操作，将EfficientNet方法放在BoTNet上：
可以看出与期望相符合，Transformer架构带来的性能上限要高于CNN，虽然模型大小比较小的时候性能比较弱，但是模型量变大以后其性能就有了保证。
代码
核心模块：MHSA （由第三方进行实现）
class MHSA(nn.Module):
def __init__(self, n_dims, width=14, height=14, heads=4):
super(MHSA, self).__init__()
self.heads = heads
self.query = nn.Conv2d(n_dims, n_dims, kernel_size=1)
self.key = nn.Conv2d(n_dims, n_dims, kernel_size=1)
self.value = nn.Conv2d(n_dims, n_dims, kernel_size=1)
self.rel_h = nn.Parameter(torch.randn([1, heads, n_dims // heads, 1, height]), requires_grad=True)
self.rel_w = nn.Parameter(torch.randn([1, heads, n_dims // heads, width, 1]), requires_grad=True)
self.softmax = nn.Softmax(dim=-1)
def forward(self, x):
n_batch, C, width, height = x.size()
q = self.query(x).view(n_batch, self.heads, C // self.heads, -1)
k = self.key(x).view(n_batch, self.heads, C // self.heads, -1)
v = self.value(x).view(n_batch, self.heads, C // self.heads, -1)
content_content = torch.matmul(q.permute(0, 1, 3, 2), k)
content_position = (self.rel_h + self.rel_w).view(1, self.heads, C // self.heads, -1).permute(0, 1, 3, 2)
content_position = torch.matmul(content_position, q)
energy = content_content + content_position
attention = self.softmax(energy)
out = torch.matmul(v, attention.permute(0, 1, 3, 2))
out = out.view(n_batch, C, width, height)
return out
参考
https://arxiv.org/abs/2101.11605
https://zhuanlan.zhihu.com/p/347849929
https://github.com/leaderj1001/BottleneckTransformers/blob/main/model.py
跑不动ImageNet，想试试Vision Transformer的同学可以看看这个仓库，
https://github.com/pprp/pytorch-cifar-model-zoo
在CIFAR10上测试：
python train.py --model 'botnet' --name "fast_training" --sched 'cosine' --epochs 100 --cutout True --lr 0.1 --bs 128 --nw 4
目前可以在100个epoch内达到验证集91.1%的准确率。