【Transformers论文】GlobalContextVisionTransformers

论文地址&＃xff1a;https://arxiv.org/abs/2206.09959https://arxiv.org/abs/2206.09959

代码地址&＃xff1a;

GitHub - NVlabs/GCVit: Official PyTorch implementation of Global Context Vision TransformersOfficial PyTorch implementation of Global Context Vision Transformers - GitHub - NVlabs/GCVit: Official PyTorch implementation of Global Context Vision Transformershttps://github.com/NVlabs/GCViT 

作者提出了全局上下文Vision Transformer(GCViT)&＃xff0c;这是一种提高参数和计算利用率的新架构。提出的方法利用全局上下文自注意模块&＃xff0c;与局部自注意相结合&＃xff0c;有效地建模长期和短期空间交互&＃xff0c;而不需要昂贵的操作。

 在这项工作中&＃xff0c;作者引入Global Context&＃xff08;GC&＃xff09;ViT 网络。提出了一个由局部和全局自注意模块组成的分层ViT架构。在每个阶段&＃xff0c;作者使用修改的fused inverted residual 模块来计算全局query token。作者称之为Fused-MBConv 模块&＃xff0c;它包含来自不同图像区域的全局上下文信息。本地自注意模块负责对short-range信息进行建模&＃xff0c;而全局query token 在所有全局自注意模块之间共享&＃xff0c;以与本地key和value进行交互。

论文主要贡献&＃xff1a;

1. 一种新的分层Transformer模型&＃xff0c;称为GCViT&＃xff0c;它可以作为各种计算机视觉任务的通用主干网络&＃xff0c;如分类、检测、实例分割&＃xff1b;
2. 一种新颖而简单的设计&＃xff0c;由全局自注意和令牌生成模块组成&＃xff0c;允许通过捕获全局上下文信息来建模长期依赖关系&＃xff0c;从而消除了对高度复杂或复杂操作的需要&＃xff1b;
3. 如图1所示&＃xff0c;实验结果SOTA。

网络结构&＃xff1a;

GC ViT 结构

网络结构如图2所示。与之前的一些Transformer架构类似&＃xff0c;使用一个层次框架&＃xff0c;通过减少空间维度&＃xff0c;同时扩大嵌入维数&＃xff0c;分别获得几个分辨率&＃xff08;称为阶段&＃xff09;的特征表示。

首先输入图像的分辨率为H X W X 3 ,通过应用一个3×3的卷积层和适当的填充来获得重叠(overlapping)的patch。然后将patch投影到C维嵌入空间中。每个GCViT阶段都由交替的局部和全局自注意模块组成&＃xff0c;以提取空间特征。两者都在像Swin Transformer 这样的 local windows 中运行&＃xff0c;然而&＃xff0c;全局自注意访问 Global Toke Generator (GTG)提取的全局特征。GTG是一个类似于cnn的模块&＃xff0c;它在每个阶段只从整个图像中提取一次特征。每个阶段后增加一个下采样模块&＃xff0c;空间分辨率减少一半。生成的特征通过平均池化和线性层传递&＃xff0c;以为下游任务创建嵌入。

从CNN模型中借用了空间特征收缩的概念&＃xff0c;该模型在降维的同时施加了局部性偏差和跨通道通信。 

class ReduceSize(nn.Module):def __init__(self, dim,norm_layer&＃61;nn.LayerNorm,keep_dim&＃61;False):super().__init__()self.conv &＃61; nn.Sequential(nn.Conv2d(dim, dim, 3, 1, 1,groups&＃61;dim, bias&＃61;False),nn.GELU(),SE(dim, dim),nn.Conv2d(dim, dim, 1, 1, 0, bias&＃61;False),)if keep_dim:dim_out &＃61; dimelse:dim_out &＃61; 2*dimself.reduction &＃61; nn.Conv2d(dim, dim_out, 3, 2, 1, bias&＃61;False)self.norm2 &＃61; norm_layer(dim_out)self.norm1 &＃61; norm_layer(dim)def forward(self, x):x &＃61; x.contiguous()x &＃61; self.norm1(x)x &＃61; x.permute(0, 3, 1, 2)x &＃61; x &＃43; self.conv(x)x &＃61; self.reduction(x).permute(0, 2, 3, 1)x &＃61; self.norm2(x)return x

Global Query Generator

作者提出包含跨整个输入特征图的信息的全局查询标记&＃xff08;global query tokens&＃xff09;&＃xff0c;以便与局部键keys和值values特征进行交互。如图5所示&＃xff1a;

class FeatExtract(nn.Module):def __init__(self, dim, keep_dim&＃61;False):super().__init__()self.conv &＃61; nn.Sequential(nn.Conv2d(dim, dim, 3, 1, 1,groups&＃61;dim, bias&＃61;False),nn.GELU(),SE(dim, dim),nn.Conv2d(dim, dim, 1, 1, 0, bias&＃61;False),)if not keep_dim:self.pool &＃61; nn.MaxPool2d(kernel_size&＃61;3, stride&＃61;2, padding&＃61;1)self.keep_dim &＃61; keep_dimdef forward(self, x):x &＃61; x.contiguous()x &＃61; x &＃43; self.conv(x)if not self.keep_dim:x &＃61; self.pool(x)return x

 Local MSA&＃xff1a;

class WindowAttention(nn.Module):def __init__(self,dim,num_heads,window_size,qkv_bias&＃61;True,qk_scale&＃61;None,attn_drop&＃61;0.,proj_drop&＃61;0.,):super().__init__()window_size &＃61; (window_size,window_size)self.window_size &＃61; window_sizeself.num_heads &＃61; num_headshead_dim &＃61; dim // num_headsself.scale &＃61; qk_scale or head_dim ** -0.5self.relative_position_bias_table &＃61; nn.Parameter(torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads))coords_h &＃61; torch.arange(self.window_size[0])coords_w &＃61; torch.arange(self.window_size[1])coords &＃61; torch.stack(torch.meshgrid([coords_h, coords_w]))coords_flatten &＃61; torch.flatten(coords, 1)relative_coords &＃61; coords_flatten[:, :, None] - coords_flatten[:, None, :]relative_coords &＃61; relative_coords.permute(1, 2, 0).contiguous()relative_coords[:, :, 0] &＃43;&＃61; self.window_size[0] - 1relative_coords[:, :, 1] &＃43;&＃61; self.window_size[1] - 1relative_coords[:, :, 0] *&＃61; 2 * self.window_size[1] - 1relative_position_index &＃61; relative_coords.sum(-1)self.register_buffer("relative_position_index", relative_position_index)self.qkv &＃61; nn.Linear(dim, dim * 3, bias&＃61;qkv_bias)self.attn_drop &＃61; nn.Dropout(attn_drop)self.proj &＃61; nn.Linear(dim, dim)self.proj_drop &＃61; nn.Dropout(proj_drop)trunc_normal_(self.relative_position_bias_table, std&＃61;.02)self.softmax &＃61; nn.Softmax(dim&＃61;-1)def forward(self, x, q_global):B_, N, C &＃61; x.shapeqkv &＃61; self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)q, k, v &＃61; qkv[0], qkv[1], qkv[2]q &＃61; q * self.scaleattn &＃61; (q &＃64; k.transpose(-2, -1))relative_position_bias &＃61; self.relative_position_bias_table[self.relative_position_index.view(-1)].view(self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1)relative_position_bias &＃61; relative_position_bias.permute(2, 0, 1).contiguous()attn &＃61; attn &＃43; relative_position_bias.unsqueeze(0)attn &＃61; self.softmax(attn)attn &＃61; self.attn_drop(attn)x &＃61; (attn &＃64; v).transpose(1, 2).reshape(B_, N, C)x &＃61; self.proj(x)x &＃61; self.proj_drop(x)return

Global MSA&＃xff1a;

class WindowAttentionGlobal(nn.Module):def __init__(self, dim,num_heads,window_size,qkv_bias&＃61;True,qk_scale&＃61;None,attn_drop&＃61;0.,proj_drop&＃61;0.,):super().__init__()window_size &＃61; (window_size,window_size)self.window_size &＃61; window_sizeself.num_heads &＃61; num_headshead_dim &＃61; dim // num_headsself.scale &＃61; qk_scale or head_dim ** -0.5self.relative_position_bias_table &＃61; nn.Parameter(torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads))coords_h &＃61; torch.arange(self.window_size[0])coords_w &＃61; torch.arange(self.window_size[1])coords &＃61; torch.stack(torch.meshgrid([coords_h, coords_w]))coords_flatten &＃61; torch.flatten(coords, 1)relative_coords &＃61; coords_flatten[:, :, None] - coords_flatten[:, None, :]relative_coords &＃61; relative_coords.permute(1, 2, 0).contiguous()relative_coords[:, :, 0] &＃43;&＃61; self.window_size[0] - 1relative_coords[:, :, 1] &＃43;&＃61; self.window_size[1] - 1relative_coords[:, :, 0] *&＃61; 2 * self.window_size[1] - 1relative_position_index &＃61; relative_coords.sum(-1)self.register_buffer("relative_position_index", relative_position_index)self.qkv &＃61; nn.Linear(dim, dim * 2, bias&＃61;qkv_bias)self.attn_drop &＃61; nn.Dropout(attn_drop)self.proj &＃61; nn.Linear(dim, dim)self.proj_drop &＃61; nn.Dropout(proj_drop)trunc_normal_(self.relative_position_bias_table, std&＃61;.02)self.softmax &＃61; nn.Softmax(dim&＃61;-1)def forward(self, x, q_global):B_, N, C &＃61; x.shapeB &＃61; q_global.shape[0]kv &＃61; self.qkv(x).reshape(B_, N, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)k, v &＃61; kv[0], kv[1]q_global &＃61; q_global.repeat(B_//B, 1, 1, 1)q &＃61; q_global.reshape(B_, self.num_heads, N, C // self.num_heads)q &＃61; q * self.scaleattn &＃61; (q &＃64; k.transpose(-2, -1))relative_position_bias &＃61; self.relative_position_bias_table[self.relative_position_index.view(-1)].view(self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1)relative_position_bias &＃61; relative_position_bias.permute(2, 0, 1).contiguous()attn &＃61; attn &＃43; relative_position_bias.unsqueeze(0)attn &＃61; self.softmax(attn)attn &＃61; self.attn_drop(attn)x &＃61; (attn &＃64; v).transpose(1, 2).reshape(B_, N, C)x &＃61; self.proj(x)x &＃61; self.proj_drop(x)return x


几种模型结构配置&＃xff1a;

实验结果&＃xff1a;