Shunted Self-Attention via Multi-Scale Token Aggregation
文章链接：https://arxiv.org/pdf/2111.15193.pdf
代码：https://github.com/OliverRensu/Shunted-Transformer
最近的Vision Transformer (ViT)模型在各种计算机视觉任务中表现优异，这得益于它能够通过自注意对图像块或标记的长期依赖进行建模。然而，这些模型通常在每一层中每个token特征指定指定的感受野。这种约束不可避免地限制了每个自注意力层对多尺度特征的捕获能力，从而导致对不同尺度多目标图像的处理性能下降。为了解决这个问题，作者提出了一种新的通用策略，称为分流自我注意(SSA)，它允许VIT在每个注意里层的混合尺度上建模注意力。SSA的关键思想是将多样的感受野传入tokens:在计算自注意力矩阵之前，它有选择地合并tokens以表示更大的对象特征，同时保留某些tokens以保留细粒度特征。这种新型的合并方案使自注意力学习不同大小对象之间的关系，同时降低了tokens数量和计算成本。通过对不同任务的大量实验验证了该算法的优越性。
一、Introduction
与专注于局部建模的卷积神经网络不同，ViTs将输入图像划分为一系列patch (token)，并通过全局自注意力的方式逐步更新token特征。自注意力可以有效地模拟token的长期依赖关系，并通过聚合来自其他token的信息逐步扩大其接受域的大小，这在很大程度上解释了VIT的成功。
然而，自注意力机制也带来了昂贵的内存消耗成本，它是输入token数量的二次多项式。因此，最先进的Transformer模型采用了各种降采样策略来减少特征大小和内存消耗。有些模型可以将第一层进行16 *16下采样投影，并对得到的粗粒度和单尺度特征图进行自注意力计算;因此，由此产生的特征信息丢失必然会降低模型的性能。其他方法努力在高分辨率特征上计算自注意力，并通过合并token与token空间缩减来降低成本。然而，这些方法倾向于在一个自注意力层中合并太多的标记，从而导致小对象标记和背景噪声的混合。这样的行为反过来又会使模型在捕捉小对象时效率降低。此外，以前的Transformer模型在很大程度上忽略了场景对象在注意力层中的多尺度性质，使得它们在涉及不同大小的对象的能力变弱。从技术上讲，这种现象归因于其潜在的注意力机制:现有方法仅依赖于标记的静态感受野和一个注意力层内统一的信息尺度，因此无法同时捕捉不同规模的特征。
为了解决这一限制，作者引入了一种新的通用自注意力方案，称为分流自我注意(SSA)，它显式地允许同一层中的自注意力头分别考虑粗粒度和细粒度特征。与之前合并太多token或无法捕获小对象的方法不同，SSA有效地在同一层内不同的注意头同时建模各种规模的对象，使自己具有良好的计算效率和保存细粒度细节的能力。
在下图中展示了普通自注意力(来自ViT)、下采样辅助注意力(来自PVT)和拟议的SSA之间的定性比较。当对相同大小的特征图应用不同的注意时，ViT捕获细粒度的小对象极高的计算成本(下图(a));PVT降低了计算成本，但它的注意力仅限于粗粒度的较大对象(b)。相比之下，SSA保持了较轻的计算负荷，但同时考虑了混合尺度的注意（c)。SSA不仅能精确地注意到粗粒度的大物体(如沙发)，还能精确地注意到细粒度的小物体(如灯和风扇)，甚至包括一些位于角落的物体。
SSA的多尺度注意机制是通过将多个注意头分成若干组来实现的。每个组都有一个专门的注意力粒度。对于细粒度组，SSA学会聚合少量token并保留更多细节。对于剩余的粗粒度头部组，SSA学习聚合大量token，从而减少计算成本，同时保留捕获大对象的能力。多粒度组共同学习多粒度信息，使模型能够有效地对多尺度对象进行建模。
总而言之，我们的贡献如下：
1、提出了分流自注意(shunt Self-Attention, SSA)，它通过多尺度标记聚合将多尺度特征提取统一在一个自注意力层内。SSA自适应地合并大对象上的标记，以提高计算效率，并保留小对象上的标记。
2、该Shunted Transformer能够有效地捕获多尺度目标，特别是小的和远程隔离的目标。
二、实现细节
（一）、Self-Attention in CNNs
卷积层的感受野通常是小而固定的。虽然膨胀卷积会扩大接受域，而变形卷积对内核有一定的偏移，但它们难以自适应和灵活地扩展到整个特征映射。受NLP领域首创的transformers自注意力层的启发，一些作品引入自注意力或非局部块来增强计算机视觉领域的卷积神经网络。这样的注意多应用在特征图尺寸较小的深层，通过多层卷积进行预处理。因此，它们不会带来太多的额外计算成本，但对性能的提升有限。
（二）、Vision Transformer
Vision Transformer (ViT)模型直接在中间层中应用自注意力，构建了无卷积神经网络模型。自开创性的ViT模型以来，许多后续工作被开发出来，通过更复杂的数据增强或知识蒸馏来提高模型的分类性能。由于自注意力的计算复杂度是token数量的二次多项式，很难直接应用于大量token。因此，这些ViT模型通常将映像划分为非重叠的大token。但是这样的划分太粗糙，会丢失很多细粒度的信息。为了保存细粒度的特征，这些模型通常向下采样特征图，并在低分辨率的特征上操作。然而，阻碍了它们在密集预测任务(如分割和检测)中的部署。
（三）、Efficient ViT Variants
为了使self-attention attention适用于大型特征映射，最近的研究提出了两种求解策略来降低计算成本:(1)将特征映射划分为区域并在区域内进行局部自注意力或(2)合并token以减少token数量。局部注意力的代表工作是Swin Transformer，它将特征映射划分为不重叠的区域，并进行局部的自注意力。然而，要通过自注意力对全局依赖进行建模，这些局部注意力需要在图像上移动窗口或堆叠很多层来获得全局感受野。对于token合并策略，PVT (Pyramid Vision Transformer)设计了一个对键与查询token合并的空间缩减关注。然而，PVT和类似的模型往往会在这种空间缩减中合并太多的token。这使得小对象的细粒度信息与背景混合，影响了模型的性能。因此，作者建议使用分流式自关注力，它可以同时保留粗粒度和细粒度的细节，同时维护对图像标记的全局依赖建模。
（四）、实现方法
本文提出的Shunted Transformer的总体架构如下图所示。它建立在新的分流自注意力(SSA)块的基础上。SSA块与传统的自注意力块在ViT中有两个主要的区别:
1)、SSA为每个自注意层引入了一个分流的注意机制，以捕获多粒度信息，更好地建模不同大小的对象，特别是小对象;
2)、通过增加交叉token交互，增强了点向前馈层提取局部信息的能力。此外，SSA部署了一种新的补丁嵌入方法，以获得更好的第一个注意块的输入特征映射。
（五）、Shunted Transformer Block
Shunted Transformer的第i级，有 L i L_i Li变压器组。每个Transformer块包含一个自注意力层和一个前馈层。为了减少处理高分辨率特征图时的计算量，PVT引入空间减少注意(spatialreduction attention, SRA)来取代原来的多头自注意力(multihead self-attention, MSA)。然而，SRA倾向于在一个自注意力层中合并太多的token，并且只提供单一规模的token特性。这些限制限制了模型捕捉多尺度物体特别是小尺度物体的能力。在一个自注意力层中并行地引入了学习多粒度的分流自注意力。
（六）、Shunted Self-Attention
首先将输入序列 F ∈ R h ∗ w ∗ c F\in R^{h*w*c} F∈Rh∗w∗c投影为查询(Q)、键(K)和值(V)。然后多头自注意力采用独立注意头H来进行并行计算自注意力。为了降低计算成本，我们遵循PVT，减小K和V的长度，而不是像Swin Transformer那样将{Q, K, V}分割成区域。
如下图所示，本文的SSA与PVT的SRA不同，在同一自注意力层的注意头上，K、V的长度不相同。相反，为了捕捉不同粒度的信息，不同头部的长度是不同的。这就产生了多尺度token聚合(MTA)。具体来说，对于i索引的不同头部，将键K和值V下采样到不同的大小。
这里的MAT ( ⋅ ; r i ) (· ; r_i) (⋅;ri)是第i个头部的多尺度token聚合层，其下采样率为 r i r_i ri。在实践中，我们取一个具有卷积核大小和步长 r i r_i ri的卷积层来实现降采样。 w i Q , W i K , W i V w_i^Q,W^K_i,W^V_i wiQ,WiK,WiV是第i个头的线性投影的参数。在注意头的一层中有不同的 r i r_i ri。因此，键和值可以捕获不同尺度的自注意力。 L E ( ⋅ ) LE(·) LE(⋅)是MTA对V值进行深度卷积的局部增强分量。与空间缩减相比，保留了更多细粒度和低级细节。那么被分流的自我注意的计算公式为:
其中 d h d_h dh是维数。由于多尺度键和值，shunted self-attention在捕捉多尺度对象时更加强大。降低的计算成本可能取决于r的值，因此我们可以很好地定义模型和r来权衡计算成本和模型性能。当r变大时，K, V中更多的token被合并，K, V的长度变短，因此计算成本低，但仍然保留了捕捉大对象的能力。而当r变小时，虽然保留了更多的细节，但带来了更多的计算成本。将各种r集成到一个自我注意层中使它能够捕获多粒度特性。
(七)、Detail-specific Feedforward Layers
在传统的前馈层中，全连接层是逐点的，不能学习交叉token信息。在这里，作者的目标是通过指定前馈层的细节来补充局部信息。如下图所示，作者通过在前馈层中的两个全连接层之间添加特定于数据的层来补充前馈层中的局部细节。
D S ( ⋅ ; θ ) DS(·;θ) DS(⋅;θ)是参数为θ的细节特定层，在实践中通过深度卷积实现。
（八）、Patch Embedding
Transformer首先是为处理顺序数据而设计的。如何将图像映射到序列中是影响模型性能的重要因素。ViT直接将输入图像分割成16个不重叠的小块。最近的一项研究发现，在补丁嵌入中使用卷积可以提供更高质量的token序列，并帮助Transformer比传统的大步幅非重叠补丁嵌入更好地观察。因此有文献采用了类似于77卷积的重叠补丁嵌入方法。
在我们的模型中，根据模型的大小取不同的卷积层进行重叠。在补丁嵌入中，以步幅为2，填充为0的77卷积层作为第一层，并根据模型大小额外增加了3个步幅为1的3*3卷积层。最后，通过一个步幅为2的非重叠投影层生成大小为 ( H / 4 ) ∗ ( W / 4 ) (H/4)*(W/4) (H/4)∗(W/4)的输入序列。
（九）、Architecture Details and Variants
给定一个大小为 H ∗ W ∗ 3 H*W*3 H∗W∗3的输入图像，采用上述补丁嵌入方案来获得长度为 ( H / 4 ) ∗ ( W / 4 ) (H/4)*(W/4) (H/4)∗(W/4)，token维数为c的信息更丰富的token序列。在之前的设计中，我们的模型分为四个阶段，每个阶段包含几个Shunted Transformer块。在每个阶段，每个块输出相同大小的特征图。我们用一个stride为2(线性嵌入)的卷积层来连接不同的阶段，特征图的大小在输入到下一阶段之前会减半，但维数会加倍。因此，我们有四个特征图F1, F2, F3, F4的每个阶段的输出，Fi的大小为 ( H / 2 i + 1 ) ∗ ( W / 2 i + 1 ) ( c ∗ 2 i − 1 ) (H/2^{i+1})*(W/2^{i+1}) (c*2^{i-1}) (H/2i+1)∗(W/2i+1)(c∗2i−1)。
如下表所示，head和 N i N_i Ni表示一个区块中head的数量和一个阶段中区块的数量。变量仅来自于不同阶段的层数。具体来说，每个区块的头数设置为2、4、8、16。patch embedding的卷积范围是1到3。
三、实验对比
为了评估Shunted Transformer的有效性，将模型应用于ImageNet-1K分类、COCO对象检测和实例分割、ADE20K语义分割。此外，还通过消融研究来评估模型中不同成分的作用。
(一)、Image Classification on ImageNet-1K
（二）、 Object Detection and Instance Segmentation
利用Mask R-CNN对COCO进行目标检测和实例分割。只有3个训练周期有多尺度训练。所有主干都在ImageNet-1K上进行预训练。*表示方法未经过同行评审。
基于COCO的RetinaNet对象检测。只有3个训练周期有多尺度训练。所有主干都在ImageNet-1K上进行预训练。*表示方法未经过同行评审。
（三）、Semantic Segmentation on ADE20K
ADE20K语义FPN和UpperNet框架中不同主干的分割性能比较.
Ade20K上Segformer框架中不同骨干的比较
不同token聚合函数在ImageNet上的精度为Top-1。
(四)、Ablation Studies
（1）、Patch Embedding
不同Patch Embedding头在ImageNet上的Top-1精度。本文的Patch Embedding需要稍微多一点的计算成本，但性能提高是显著的。
（2）、Token Aggregation Function
不同token聚合函数在ImageNet上Top-1的精度
(3)、Detail-specific Feed-Forward
通过类似的参数数和FLOPs，具体细节的前馈层在ImageNet上比传统前馈层提供更高的top-1精度。
在本文中，作者提出了一种新的分流自注意(SSA)方案，以显式地考虑多尺度特征。与以往只关注一个注意层中的静态特征图不同，本文在一个自注意力层中维护了关注多尺度物体的各种尺度特征图。大量的实验表明，该模型作为各种下游任务的主干是有效的。
源码如下：
import torch
import torch.nn as nn
from functools import partial
from timm.models.layers import DropPath, to_2tuple, trunc_normal_
import math
class Mlp(nn.Module):
def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
super().__init__()
out_features = out_features or in_features
hidden_features = hidden_features or in_features
self.fc1 = nn.Linear(in_features, hidden_features)
self.dwconv = DWConv(hidden_features)
self.act = act_layer()
self.fc2 = nn.Linear(hidden_features, out_features)
self.drop = nn.Dropout(drop)
self.apply(self._init_weights)
def _init_weights(self, m):
if isinstance(m, nn.Linear):
trunc_normal_(m.weight, std=.02)
if isinstance(m, nn.Linear) and m.bias is not None:
nn.init.constant_(m.bias, 0)
elif isinstance(m, nn.LayerNorm):
nn.init.constant_(m.bias, 0)
nn.init.constant_(m.weight, 1.0)
elif isinstance(m, nn.Conv2d):
fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
fan_out //= m.groups
m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))
if m.bias is not None:
m.bias.data.zero_()
def forward(self, x, H, W):
x = self.fc1(x)
x = self.act(x + self.dwconv(x, H, W))
x = self.drop(x)
x = self.fc2(x)
x = self.drop(x)
return x
class Attention(nn.Module):
def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0., sr_ratio=1):
super().__init__()
assert dim % num_heads == 0, f"dim {dim} should be divided by num_heads {num_heads}."
self.dim = dim
self.num_heads = num_heads
head_dim = dim // num_heads
self.scale = qk_scale or head_dim ** -0.5
self.q = nn.Linear(dim, dim, bias=qkv_bias)
self.attn_drop = nn.Dropout(attn_drop)
self.proj = nn.Linear(dim, dim)
self.proj_drop = nn.Dropout(proj_drop)
self.sr_ratio = sr_ratio
if sr_ratio > 1:
self.act = nn.GELU()
if sr_ratio == 8:
self.sr1 = nn.Conv2d(dim, dim, kernel_size=8, stride=8)
self.norm1 = nn.LayerNorm(dim)
self.sr2 = nn.Conv2d(dim, dim, kernel_size=4, stride=4)
self.norm2 = nn.LayerNorm(dim)
if sr_ratio == 4:
self.sr1 = nn.Conv2d(dim, dim, kernel_size=4, stride=4)
self.norm1 = nn.LayerNorm(dim)
self.sr2 = nn.Conv2d(dim, dim, kernel_size=2, stride=2)
self.norm2 = nn.LayerNorm(dim)
if sr_ratio == 2:
self.sr1 = nn.Conv2d(dim, dim, kernel_size=2, stride=2)
self.norm1 = nn.LayerNorm(dim)
self.sr2 = nn.Conv2d(dim, dim, kernel_size=1, stride=1)
self.norm2 = nn.LayerNorm(dim)
self.kv1 = nn.Linear(dim, dim, bias=qkv_bias)
self.kv2 = nn.Linear(dim, dim, bias=qkv_bias)
self.local_conv1 = nn.Conv2d(dim//2, dim//2, kernel_size=3, padding=1, stride=1, groups=dim//2)
self.local_conv2 = nn.Conv2d(dim//2, dim//2, kernel_size=3, padding=1, stride=1, groups=dim//2)
else:
self.kv = nn.Linear(dim, dim * 2, bias=qkv_bias)
self.local_conv = nn.Conv2d(dim, dim, kernel_size=3, padding=1, stride=1, groups=dim)
self.apply(self._init_weights)
def _init_weights(self, m):
if isinstance(m, nn.Linear):
trunc_normal_(m.weight, std=.02)
if isinstance(m, nn.Linear) and m.bias is not None:
nn.init.constant_(m.bias, 0)
elif isinstance(m, nn.LayerNorm):
nn.init.constant_(m.bias, 0)
nn.init.constant_(m.weight, 1.0)
elif isinstance(m, nn.Conv2d):
fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
fan_out //= m.groups
m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))
if m.bias is not None:
m.bias.data.zero_()
def forward(self, x, H, W):
B, N, C = x.shape
q = self.q(x).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
if self.sr_ratio > 1:
x_ = x.permute(0, 2, 1).reshape(B, C, H, W)
x_1 = self.act(self.norm1(self.sr1(x_).reshape(B, C, -1).permute(0, 2, 1)))
x_2 = self.act(self.norm2(self.sr2(x_).reshape(B, C, -1).permute(0, 2, 1)))
kv1 = self.kv1(x_1).reshape(B, -1, 2, self.num_heads//2, C // self.num_heads).permute(2, 0, 3, 1, 4)
kv2 = self.kv2(x_2).reshape(B, -1, 2, self.num_heads//2, C // self.num_heads).permute(2, 0, 3, 1, 4)
k1, v1 = kv1[0], kv1[1] #B head N C
k2, v2 = kv2[0], kv2[1]
attn1 = (q[:, :self.num_heads//2] @ k1.transpose(-2, -1)) * self.scale
attn1 = attn1.softmax(dim=-1)
attn1 = self.attn_drop(attn1)
v1 = v1 + self.local_conv1(v1.transpose(1, 2).reshape(B, -1, C//2).
transpose(1, 2).view(B, C//2, H//self.sr_ratio, W//self.sr_ratio)).\
view(B, C//2, -1).view(B, self.num_heads//2, C // self.num_heads, -1).transpose(-1, -2)
x1 = (attn1 @ v1).transpose(1, 2).reshape(B, N, C//2)
attn2 = (q[:, self.num_heads // 2:] @ k2.transpose(-2, -1)) * self.scale
attn2 = attn2.softmax(dim=-1)
attn2 = self.attn_drop(attn2)
v2 = v2 + self.local_conv2(v2.transpose(1, 2).reshape(B, -1, C//2).
transpose(1, 2).view(B, C//2, H*2//self.sr_ratio, W*2//self.sr_ratio)).\
view(B, C//2, -1).view(B, self.num_heads//2, C // self.num_heads, -1).transpose(-1, -2)
x2 = (attn2 @ v2).transpose(1, 2).reshape(B, N, C//2)
x = torch.cat([x1, x2], dim=-1)
else:
kv = self.kv(x).reshape(B, -1, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
k, v = kv[0], kv[1]
attn = (q @ k.transpose(-2, -1)) * self.scale
attn = attn.softmax(dim=-1)
attn = self.attn_drop(attn)
x = (attn @ v).transpose(1, 2).reshape(B, N, C) + self.local_conv(v.transpose(1, 2).reshape(B, N, C).
transpose(1, 2).view(B, C, H, W)).view(B, C, N).transpose(1, 2)
x = self.proj(x)
x = self.proj_drop(x)
return x
class Block(nn.Module):
def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, sr_ratio=1):
super().__init__()
self.norm1 = norm_layer(dim)
self.attn = Attention(
dim,
num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale,
attn_drop=attn_drop, proj_drop=drop, sr_ratio=sr_ratio)
# NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
self.norm2 = norm_layer(dim)
mlp_hidden_dim = int(dim * mlp_ratio)
self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
self.apply(self._init_weights)
def _init_weights(self, m):
if isinstance(m, nn.Linear):
trunc_normal_(m.weight, std=.02)
if isinstance(m, nn.Linear) and m.bias is not None:
nn.init.constant_(m.bias, 0)
elif isinstance(m, nn.LayerNorm):
nn.init.constant_(m.bias, 0)
nn.init.constant_(m.weight, 1.0)
elif isinstance(m, nn.Conv2d):
fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
fan_out //= m.groups
m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))
if m.bias is not None:
m.bias.data.zero_()
def forward(self, x, H, W):
x = x + self.drop_path(self.attn(self.norm1(x), H, W))
x = x + self.drop_path(self.mlp(self.norm2(x), H, W))
return x
class OverlapPatchEmbed(nn.Module):
""" Image to Patch Embedding
"""
def __init__(self, img_size=224, patch_size=7, stride=4, in_chans=3, embed_dim=768):
super().__init__()
img_size = to_2tuple(img_size)
patch_size = to_2tuple(patch_size)
self.img_size = img_size
self.patch_size = patch_size
self.H, self.W = img_size[0] // patch_size[0], img_size[1] // patch_size[1]
self.num_patches = self.H * self.W
self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride,
padding=(patch_size[0] // 2, patch_size[1] // 2))
self.norm = nn.LayerNorm(embed_dim)
self.apply(self._init_weights)
def _init_weights(self, m):
if isinstance(m, nn.Linear):
trunc_normal_(m.weight, std=.02)
if isinstance(m, nn.Linear) and m.bias is not None:
nn.init.constant_(m.bias, 0)
elif isinstance(m, nn.LayerNorm):
nn.init.constant_(m.bias, 0)
nn.init.constant_(m.weight, 1.0)
elif isinstance(m, nn.Conv2d):
fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
fan_out //= m.groups
m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))
if m.bias is not None:
m.bias.data.zero_()
def forward(self, x):
x = self.proj(x)
_, _, H, W = x.shape
x = x.flatten(2).transpose(1, 2)
x = self.norm(x)
return x, H, W
class ShuntedTransformer(nn.Module):
def __init__(self, cin=512, depth=3, img_size=40,
num_heads=2, mlp_ratios=4, qkv_bias=True, qk_scale=None, drop_rate=0.,
attn_drop_rate=0., drop_path_rate=0., norm_layer=nn.LayerNorm,
sr_ratios=2):
super().__init__()
self.depth = depth
dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]
cur = 0
self.patch_embed = OverlapPatchEmbed(img_size=img_size, patch_size=3, stride=1, in_chans=cin, embed_dim=cin)
self.block = nn.ModuleList([Block(
dim=cin, num_heads=num_heads, mlp_ratio=mlp_ratios, qkv_bias=qkv_bias, qk_scale=qk_scale,
drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + j], norm_layer=norm_layer,
sr_ratio=sr_ratios)
for j in range(depth)])
self.norm = norm_layer(cin)
cur += depth
self.apply(self._init_weights)
def _init_weights(self, m):
if isinstance(m, nn.Linear):
trunc_normal_(m.weight, std=.02)
if isinstance(m, nn.Linear) and m.bias is not None:
nn.init.constant_(m.bias, 0)
elif isinstance(m, nn.LayerNorm):
nn.init.constant_(m.bias, 0)
nn.init.constant_(m.weight, 1.0)
elif isinstance(m, nn.Conv2d):
fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
fan_out //= m.groups
m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))
if m.bias is not None:
m.bias.data.zero_()
def freeze_patch_emb(self):
self.patch_embed1.requires_grad = False
@torch.jit.ignore
def no_weight_decay(self):
return {'pos_embed1', 'pos_embed2', 'pos_embed3', 'pos_embed4', 'cls_token'} # has pos_embed may be better
def forward_features(self, x):
B = x.shape[0]
x, H, W = self.patch_embed(x)
for blk in self.block:
x = blk(x, H, W)
x = self.norm(x)
x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()
return x
def forward(self, x):
x = self.forward_features(x)
return x
class DWConv(nn.Module):
def __init__(self, dim=768):
super(DWConv, self).__init__()
self.dwconv = nn.Conv2d(dim, dim, 3, 1, 1, bias=True, groups=dim)
def forward(self, x, H, W):
B, N, C = x.shape
x = x.transpose(1, 2).view(B, C, H, W)
x = self.dwconv(x)
x = x.flatten(2).transpose(1, 2)
return x
def _conv_filter(state_dict, patch_size=16):
""" convert patch embedding weight from manual patchify + linear proj to conv"""
out_dict = {}
for k, v in state_dict.items():
if 'patch_embed.proj.weight' in k:
v = v.reshape((v.shape[0], 3, patch_size, patch_size))
out_dict[k] = v
return out_dict
if __name__ == '__main__':
model = ShuntedTransformer(
cin=512, qkv_bias=True,
norm_layer=partial(nn.LayerNorm, eps=1e-6), depth=3)
input = torch.randn(1, 512, 40, 40) # 模型输入维数
x = model(input)
print(x.shape)