image.png

主要工作

  1. 结构上的设计
    1. 使用卷积tokenizer来从图像中提取token特征，空间下采样1/4
    2. tokenization后使用一个纯卷积结构再下采样1/2
    3. 替换空间mlp为通道mlp和空间卷积的组合，这可以打破对于输入维度的限定（卷积的优势）
    4. 引入多尺度结构，使用跨步卷积实现下采样
  2. 卷积的引入了补充了空间上的交互

  3. 整体结构具有更少的参数量（这主要还是因为以上卷积策略的应用）

    关键图示

    image.png
    image.png

    核心代码

    1. class ConvMLPStage(Module):
    2.  def __init__(self,
    3.               embedding_dim,
    4.               dim_feedforward=2048,
    5.               stochastic_depth_rate=0.1):
    6.      super(ConvMLPStage, self).__init__()
    7.      self.norm1 = LayerNorm(embedding_dim)
    8.      self.channel_mlp1 = Mlp(embedding_dim_in=embedding_dim, hidden_dim=dim_feedforward)
    9.      self.norm2 = LayerNorm(embedding_dim)
    10.      self.connect = Conv2d(embedding_dim,
    11.                            embedding_dim,
    12.                            kernel_size=(3, 3),
    13.                            stride=(1, 1),
    14.                            padding=(1, 1),
    15.                            groups=embedding_dim,
    16.                            bias=False)
    17.      self.connect_norm = LayerNorm(embedding_dim)
    18.      self.channel_mlp2 = Mlp(embedding_dim_in=embedding_dim, hidden_dim=dim_feedforward)
    19.      self.drop_path = DropPath(stochastic_depth_rate) if stochastic_depth_rate > 0 else Identity()
    21.  def forward(self, src):
    22.      src = src + self.drop_path(self.channel_mlp1(self.norm1(src)))
    23.      src = self.connect(self.connect_norm(src).permute(0, 3, 1, 2)).permute(0, 2, 3, 1)
    24.      src = src + self.drop_path(self.channel_mlp2(self.norm2(src)))
    25.      return src

    链接