CoTNet-重磅开源！京东AI Research提出新的主干网络CoTNet,在CVPR上获得开放域图像识别竞赛冠军

基于Self-Attention的Transformer结构，首先在NLP任务中被提出，最近在CV任务中展现出了非常好的效果。然而，大多数现有的Transformer直接在二维特征图上的进行Self-Attention，基于每个空间位置的query和key获得注意力矩阵，但相邻的key之间的上下文信息未得到充分利用。

本文设计了一种新的注意力结构CoT Block，这种结构充分利用了key的上下文信息，以指导动态注意力矩阵的学习，从而增强了视觉表示的能力。

作者将CoT Block代替了ResNet结构中的3x3卷积，来形成CoTNet，最终在一系列视觉任务（分类、检测、分割）上取得了非常好的性能，此外，CoTNet在CVPR上获得开放域图像识别竞赛冠军。

import numpy as np
import torch
from torch import flatten, nn
from torch.nn import init
from torch.nn.modules.activation import ReLU
from torch.nn.modules.batchnorm import BatchNorm2d
from torch.nn import functional as F



class CoTAttention(nn.Module):

    def __init__(self, dim=512,kernel_size=3):
        super().__init__()
        self.dim=dim
        self.kernel_size=kernel_size

        self.key_embed=nn.Sequential(
            nn.Conv2d(dim,dim,kernel_size=kernel_size,padding=kernel_size//2,groups=4,bias=False),
            nn.BatchNorm2d(dim),
            nn.ReLU()
        )
        self.value_embed=nn.Sequential(
            nn.Conv2d(dim,dim,1,bias=False),
            nn.BatchNorm2d(dim)
        )

        factor=4
        self.attention_embed=nn.Sequential(
            nn.Conv2d(2*dim,2*dim//factor,1,bias=False),
            nn.BatchNorm2d(2*dim//factor),
            nn.ReLU(),
            nn.Conv2d(2*dim//factor,kernel_size*kernel_size*dim,1)
        )


    def forward(self, x):
        bs,c,h,w=x.shape
        k1=self.key_embed(x) #bs,c,h,w
        v=self.value_embed(x).view(bs,c,-1) #bs,c,h,w

        y=torch.cat([k1,x],dim=1) #bs,2c,h,w
        att=self.attention_embed(y) #bs,c*k*k,h,w
        att=att.reshape(bs,c,self.kernel_size*self.kernel_size,h,w)
        att=att.mean(2,keepdim=False).view(bs,c,-1) #bs,c,h*w
        k2=F.softmax(att,dim=-1)*v
        k2=k2.view(bs,c,h,w)


        return k1+k2


if __name__ == '__main__':
    input=torch.randn(50,512,7,7)
    cot = CoTAttention(dim=512,kernel_size=3)
    output=cot(input)
    print(output.shape)