点击下方卡片，关注“小白玩转Python”公众号

截至2024年1月，ViTDet是所有视觉任务的首选架构。它被用于“segment-anything”。在ViTAE-Transformer中，我们在语义分割、目标检测、人体姿势、抠图、遥感等多个任务上取得了最先进的结果。理解这个骨干架构将有助于我们根据任务选择最佳参数。

ViTDet的设计是为了强调使用变换器进行目标检测的专门架构的必要性。从某种意义上说，我会将其称为一个超简化的Swin Transformers，基本上去掉了网络的分层结构，转换了窗口等。注意：我们只讨论骨干部分，不涉及基于FPN的消融研究。

因此，网络大致分为以下几部分：

[PatchEmbed] -> nx[blocks] -> [Neck]

在每个块内部，我们有：

  • 窗口注意力 

  • 相对位置编码

导入所需的参数

import math
import numpy as np
import torch
import torch.nn as nn
import fastcore.all as fc
from PIL import Image
from functools import partial
from torchvision.transforms import RandomResizedCrop, RandomHorizontalFlip, Compose, ToTensor, ToPILImage

让我们创建一个大小为224x224，patch 大小为32的图像

img_size = 1024
patch_size = 32

加载和可视化图像 

我们加载并使用coco val数据。对于本博客目的，您可以从互联网上选择任意图像。 

imgs = fc.L(fc.Path("coco/val2017/").glob("*.jpg"))
imgs
(#5000) [Path('coco/val2017/000000182611.jpg'),Path('coco/val2017/000000335177.jpg'),Path('coco/val2017/000000278705.jpg'),Path('coco/val2017/000000463618.jpg'),Path('coco/val2017/000000568981.jpg'),Path('coco/val2017/000000092416.jpg'),Path('coco/val2017/000000173830.jpg'),Path('coco/val2017/000000476215.jpg'),Path('coco/val2017/000000479126.jpg'),Path('coco/val2017/000000570664.jpg')...]

以下是将图像调整为所需形状的基本转换：

def transforms():
    return Compose([RandomResizedCrop(size=1024, scale=[0.4, 1], ratio=[0.75, 1.33], interpolation=2), 
                    RandomHorizontalFlip(p=0.5), 
                    ToTensor()])


def load_img(img_loc, transforms):
    img = Image.open(img_loc)
    return transforms(img)


load_img = partial(load_img, transforms=transforms())
img = load_img(imgs[1])
img.shape #torch.Size([3, 1024, 1024])

Patch Embed

我们将为[3x32x32]创建补丁嵌入。为此，我们可以使用一个简单的卷积层，内核和步幅均为补丁大小

num_channels = 3
hidden_size = 768
projection = nn.Conv2d(num_channels, hidden_size, kernel_size=patch_size, stride=patch_size)
projection #Conv2d(3, 768, kernel_size=(32, 32), stride=(32, 32))
pe = projection(img.unsqueeze(0))
pe.shape #torch.Size([1, 768, 32, 32])

重新排列像素：

pe = pe.permute((0, 2, 3, 1))
pe.shape #torch.Size([1, 32, 32, 768])

现在我们有了[32x32] = 1024个令牌，每个令牌有768个向量。使用卷积类型结构保留了每个令牌相对于其他令牌的位置。我们可以将位置编码添加到这些特征中（可选）。

Transformer Blocks

在每个Transformer块中，我们首先应用窗口化，然后计算注意力，重新连接窗口块，应用mlp。Transformer块还具有一些跳过连接和规范化层，如下所示。

7729cf3b382951438fc6c25e1b1bdcba.png

ViTDet中的Transformer块

窗口化

在“ViTDet”的上下文中，窗口化是可选的，可以对所有令牌进行注意力计算。这种类型的注意力称为“全局注意力”。但是，全局注意力很昂贵，因为在本例中我们必须计算一个32x32的矩阵。如下图所示，如果补丁大小要小得多，则注意力矩阵的计算量会呈二次增长，使得计算变得非常昂贵。因此，考虑到窗口化注意力。

3654cd320849a91d231b411c4f8d522b.png

不同补丁大小的注意力矩阵大小

首先，将32x32矩阵分成8x8（窗口大小）窗口。因此，我们将获得总共（32/8）*（32/8）= 16个窗口，每个窗口具有（8x8）64个令牌。只在这些令牌内计算注意力，使其成为局部注意力。

692a728930091f1be31edbf32bfac3fe.png

窗口化注意力矩阵

8c3e084a1c9e39bf5c46dd762f882f3a.png

不同补丁大小的窗口化注意力矩阵

从上述两个表中，我们可以看出窗口化注意力计算更加可行，且内存占用更少。

window_size = 8
batch_size, height, width, num_channels = pe.shape
wpe = pe.view(
        batch_size, height // window_size, window_size, width // window_size, window_size, num_channels
    )
wpe.shape #torch.Size([1, 4, 8, 4, 8, 768])
windows = wpe.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, num_channels)
windows.shape #torch.Size([16, 8, 8, 768])
windows = windows.view(-1, window_size*window_size, num_channels)
windows.shape #torch.Size([16, 64, 768])

注意力

这是一个简单的注意力，如“注意力是你所需要的”论文中所讨论的。我们将一步一步地看到如下过程：

12568665413ad908c9e601f2c091c17b.png

ViTDet中的注意力

我们通过使用MLP层获取q、k、v矩阵：

dim = windows.shape[-1]
num_heads = 4
head_dim = dim // num_heads
scale = head_dim**-0.5
wq = [nn.Linear(dim, head_dim) for head in range(num_heads)]
wk = [nn.Linear(dim, head_dim) for head in range(num_heads)]
wv = [nn.Linear(dim, head_dim) for head in range(num_heads)]
wq, wk, wv
([Linear(in_features=768, out_features=192, bias=True),
  Linear(in_features=768, out_features=192, bias=True),
  Linear(in_features=768, out_features=192, bias=True),
  Linear(in_features=768, out_features=192, bias=True)],
 [Linear(in_features=768, out_features=192, bias=True),
  Linear(in_features=768, out_features=192, bias=True),
  Linear(in_features=768, out_features=192, bias=True),
  Linear(in_features=768, out_features=192, bias=True)],
 [Linear(in_features=768, out_features=192, bias=True),
  Linear(in_features=768, out_features=192, bias=True),
  Linear(in_features=768, out_features=192, bias=True),
  Linear(in_features=768, out_features=192, bias=True)])
q = [i(windows) for i in wq]
k = [i(windows) for i in wk]
v = [i(windows) for i in wv]
[i.shape for i in q] ##
#[torch.Size([16, 64, 192]),
# torch.Size([16, 64, 192]),
# torch.Size([16, 64, 192]),
# torch.Size([16, 64, 192])]
q = torch.concatenate(q) # number of heads * windows
k = torch.concatenate(k)
v = torch.concatenate(v)
q.shape, k.shape, v.shape
(torch.Size([64, 64, 192]),
 torch.Size([64, 64, 192]),
 torch.Size([64, 64, 192]))

对q和k进行矩阵乘法并使用比例：

attention_scores = (q @ k.transpose(-2, -1)) * scale
attention_scores.shape #torch.Size([64, 64, 64])

应用相对位置编码：

这是一个独立的话题，但基本上我们会在每个注意力块中添加位置编码，而不是像在普通的Vanilla Vit中那样在开始时添加。

rel_pos_h = nn.Parameter(torch.zeros(2 * window_size - 1, head_dim))
rel_pos_w = nn.Parameter(torch.zeros(2 * window_size - 1, head_dim))
rel_pos_h.shape, rel_pos_w.shape #(torch.Size([15, 192]), torch.Size([15, 192]))
from transformers.models.vitdet.modeling_vitdet import add_decomposed_relative_positions
attention_scores = add_decomposed_relative_positions(
                attention_scores, q, rel_pos_h, rel_pos_w, (window_size, window_size), (window_size, window_size)
            )
attention_scores.shape #torch.Size([64, 64, 64])

应用softmax：

attention_probs = attention_scores.softmax(dim=-1)
attention_probs.shape #torch.Size([64, 64, 64])

乘以key向量：

hidden_state = attention_probs @ v
hidden_state.shape #torch.Size([64, 64, 192])
hidden_state = hidden_state.view(16, num_heads, window_size, window_size, -1)
hidden_state = hidden_state.permute(0, 2, 3, 1, 4)
hidden_state = hidden_state.reshape(16, window_size, window_size, -1)
hidden_state.shape #torch.Size([16, 8, 8, 768])

添加投影层：

proj = nn.Linear(dim, dim)
proj #Linear(in_features=768, out_features=768, bias=True)
attention_out = proj(hidden_state)
attention_out.shape #torch.Size([16, 8, 8, 768])

去窗口化

去窗口化现有向量，将其变为（batch_size，tokens，embedding_dim）的形式：

pe = attention_out.view(-1, height // window_size, width // window_size, \
                       window_size, window_size, num_channels)
pe = pe.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, height, width, num_channels)
pe.shape #torch.Size([1, 32, 32, 768])

我们已经进行了窗口化——应用了注意力——取消了窗口化以获得向量。我们可以看到，输出向量大小与输入大小相同。如果不想进行全局注意力，则可以将窗口大小设置为输入大小。在这种情况下，它是32x32。

残差块

到目前为止，我们已经看到注意力只在窗口内应用。为了跨窗口学习，我们确实在某些层中应用了全局注意力。全局注意力被认为是昂贵的，因此仅在少数情况下应用。

  • 网络被划分为4个子集。每个子集包含6个块。因此，总共有24个层。

  • 在每个子集的最后一个块末尾，我们应用全局注意力。这将减少我们的计算量，也允许令牌在窗口外学习。

论文的作者还建议使用卷积层的残差块代替全局注意力。网络如下所示，具有1x1、3x3和1x1的卷积层。这将允许网络从所有令牌中学习。

from transformers.models.vitdet.modeling_vitdet import VitDetResBottleneckBlock
class config:
    hidden_act = "gelu"
residual = VitDetResBottleneckBlock(config, in_channels=768, out_channels=768, bottleneck_channels=768//2)
residual
VitDetResBottleneckBlock(
  (conv1): Conv2d(768, 384, kernel_size=(1, 1), stride=(1, 1), bias=False)
  (norm1): VitDetLayerNorm()
  (act1): GELUActivation()
  (conv2): Conv2d(384, 384, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
  (norm2): VitDetLayerNorm()
  (act2): GELUActivation()
  (conv3): Conv2d(384, 768, kernel_size=(1, 1), stride=(1, 1), bias=False)
  (norm3): VitDetLayerNorm()
)
residual(pe.permute((0, 3, 1, 2))).shape #torch.Size([1, 768, 32, 32])

网络的关键部分只有这些。现在让我们定义Segment Anything骨干中的所有参数，并查看是否一切都说得通。

完整的网络结构

from segment_anything.modeling.image_encoder import ImageEncoderViT
enc = ImageEncoderViT(img_size=1024,
                      patch_size=16, 
                      in_chans=3, 
                      embed_dim=768, 
                      depth=12, 
                      num_heads=12, 
                      mlp_ratio=4, 
                      out_chans=256, 
                      qkv_bias=True, 
                      norm_layer= torch.nn.modules.normalization.LayerNorm, 
                      act_layer=torch.nn.modules.activation.GELU, 
                      use_abs_pos=False, 
                      use_rel_pos=True, 
                      window_size=16,
                      global_attn_indexes=[2, 5, 8, 11])
enc
ImageEncoderViT(
  (patch_embed): PatchEmbed(
    (proj): Conv2d(3, 768, kernel_size=(16, 16), stride=(16, 16))
  )
  (blocks): ModuleList(
    (0-11): 12 x Block(
      (norm1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
      (attn): Attention(
        (qkv): Linear(in_features=768, out_features=2304, bias=True)
        (proj): Linear(in_features=768, out_features=768, bias=True)
      )
      (norm2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
      (mlp): MLPBlock(
        (lin1): Linear(in_features=768, out_features=3072, bias=True)
        (lin2): Linear(in_features=3072, out_features=768, bias=True)
        (act): GELU(approximate='none')
      )
    )
  )
  (neck): Sequential(
    (0): Conv2d(768, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
    (1): LayerNorm2d()
    (2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
    (3): LayerNorm2d()
  )
)

消融研究

  • 在辅以少量全局注意力块时，窗口化注意力就足够了。

  • 使用残差卷积或全局注意力可以获得类似的性能。使用残差卷积时，训练和推断时间要低得多。

  • 掩码自编码器提供了强大的预训练骨干

  • 与层级骨干（如MViT2或Swin Transformers）相比，ViTDet效果更好。

  • 当使用Imagenet 1k进行预训练时，最终在coco测试集上达到了61.3 APbox。

·  END  ·

HAPPY LIFE

e7a43cb9de30056840fac075404e87a0.png

本文仅供学习交流使用，如有侵权请联系作者删除

Logo
开放原子开发者工作坊

开放原子开发者工作坊旨在鼓励更多人参与开源活动，与志同道合的开发者们相互交流开发经验、分享开发心得、获取前沿技术趋势。工作坊有多种形式的开发者活动，如meetup、训练营等，主打技术交流，干货满满，真诚地邀请各位开发者共同参与！

更多推荐

第二届开放原子大赛首批创新成果集结武汉，诚邀广大开发者共鉴开源技术盛宴

第二届开放原子大赛首批创新成果集结武汉，诚邀广大开发者共鉴开源技术盛宴

avatar 开放原子开发者工作坊

诚邀报名 | 开源基础设施能力建设分论坛：打造开源生态的“心脏”

诚邀报名 | 开源基础设施能力建设分论坛：打造开源生态的“心脏”

avatar 开放原子开发者工作坊

诚邀报名 | 编程语言分论坛：AI时代的技术革新与开源实践

诚邀报名 | 编程语言分论坛：AI时代的技术革新与开源实践

avatar 开放原子开发者工作坊