点击下方卡片，关注“小白玩转Python”公众号
截至2024年1月，ViTDet是所有视觉任务的首选架构。它被用于“segment-anything”。在ViTAE-Transformer中，我们在语义分割、目标检测、人体姿势、抠图、遥感等多个任务上取得了最先进的结果。理解这个骨干架构将有助于我们根据任务选择最佳参数。
ViTDet的设计是为了强调使用变换器进行目标检测的专门架构的必要性。从某种意义上说，我会将其称为一个超简化的Swin Transformers，基本上去掉了网络的分层结构，转换了窗口等。注意：我们只讨论骨干部分，不涉及基于FPN的消融研究。
因此，网络大致分为以下几部分：
[PatchEmbed] -> nx[blocks] -> [Neck]
在每个块内部，我们有：
导入所需的参数
import mathimport numpy as npimport torchimport torch.nn as nnimport fastcore.all as fcfrom PIL import Imagefrom functools import partialfrom torchvision.transforms import RandomResizedCrop, RandomHorizontalFlip, Compose, ToTensor, ToPILImage
让我们创建一个大小为224x224，patch 大小为32的图像
img_size = 1024patch_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
Patch Embed
我们将为[3x32x32]创建补丁嵌入。为此，我们可以使用一个简单的卷积层，内核和步幅均为补丁大小
num_channels = 3hidden_size = 768projection = nn.Conv2d(num_channels, hidden_size, kernel_size=patch_size, stride=patch_size)projection
pe = projection(img.unsqueeze(0))pe.shape
重新排列像素：
pe = pe.permute((0, 2, 3, 1))pe.shape
现在我们有了[32x32] = 1024个令牌，每个令牌有768个向量。使用卷积类型结构保留了每个令牌相对于其他令牌的位置。我们可以将位置编码添加到这些特征中（可选）。
Transformer Blocks
在每个Transformer块中，我们首先应用窗口化，然后计算注意力，重新连接窗口块，应用mlp。Transformer块还具有一些跳过连接和规范化层，如下所示。
ViTDet中的Transformer块
窗口化
在“ViTDet”的上下文中，窗口化是可选的，可以对所有令牌进行注意力计算。这种类型的注意力称为“全局注意力”。但是，全局注意力很昂贵，因为在本例中我们必须计算一个32x32的矩阵。如下图所示，如果补丁大小要小得多，则注意力矩阵的计算量会呈二次增长，使得计算变得非常昂贵。因此，考虑到窗口化注意力。
不同补丁大小的注意力矩阵大小
首先，将32x32矩阵分成8x8（窗口大小）窗口。因此，我们将获得总共（32/8）*（32/8）= 16个窗口，每个窗口具有（8x8）64个令牌。只在这些令牌内计算注意力，使其成为局部注意力。
窗口化注意力矩阵
不同补丁大小的窗口化注意力矩阵
从上述两个表中，我们可以看出窗口化注意力计算更加可行，且内存占用更少。
window_size = 8batch_size, height, width, num_channels = pe.shapewpe = pe.view( batch_size, height // window_size, window_size, width // window_size, window_size, num_channels )wpe.shape
windows = wpe.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, num_channels)windows.shape
windows = windows.view(-1, window_size*window_size, num_channels)windows.shape
注意力
这是一个简单的注意力，如“注意力是你所需要的”论文中所讨论的。我们将一步一步地看到如下过程：
ViTDet中的注意力
我们通过使用MLP层获取q、k、v矩阵：
dim = windows.shape[-1]num_heads = 4head_dim = dim // num_headsscale = head_dim**-0.5wq = [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]
q = torch.concatenate(q) 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)) * scaleattention_scores.shape
应用相对位置编码：
这是一个独立的话题，但基本上我们会在每个注意力块中添加位置编码，而不是像在普通的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
from transformers.models.vitdet.modeling_vitdet import add_decomposed_relative_positionsattention_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
应用softmax：
attention_probs = attention_scores.softmax(dim=-1)attention_probs.shape
乘以key向量：
hidden_state = attention_probs @ vhidden_state.shape
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
添加投影层：
proj = nn.Linear(dim, dim)proj
attention_out = proj(hidden_state)attention_out.shape
去窗口化
去窗口化现有向量，将其变为（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
我们已经进行了窗口化——应用了注意力——取消了窗口化以获得向量。我们可以看到，输出向量大小与输入大小相同。如果不想进行全局注意力，则可以将窗口大小设置为输入大小。在这种情况下，它是32x32。
残差块
到目前为止，我们已经看到注意力只在窗口内应用。为了跨窗口学习，我们确实在某些层中应用了全局注意力。全局注意力被认为是昂贵的，因此仅在少数情况下应用。
论文的作者还建议使用卷积层的残差块代替全局注意力。网络如下所示，具有1x1、3x3和1x1的卷积层。这将允许网络从所有令牌中学习。
from transformers.models.vitdet.modeling_vitdet import VitDetResBottleneckBlockclass config: hidden_act = "gelu"residual = VitDetResBottleneckBlock(config, in_channels=768, out_channels=768, bottleneck_channels=768residual
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。
