VitDet学习笔记

github：detectron2/projects/ViTDet at main · facebookresearch/detectron2 · GitHub

骨干网络两种Mask RCNN，cascade Mask RCNN

vitdet即vit-b Mask RCNN和MAE的衍生

vit结构：

 核心代码：多头自注意力机制

class attention(nn.Module):
    # 初始化
    def __init__(self, dim, num_heads=12, qkv_bias=False, atten_drop_ratio=0., proj_drop_ratio=0.):
        super(attention, self).__init__()

        # 多头注意力的数量
        self.num_heads = num_heads
        # 将生成的qkv均分成num_heads个。得到每个head的qkv对应的通道数。
        head_dim = dim // num_heads
        # 公式中的分母
        self.scale = head_dim ** -0.5
        # 通过一个全连接层计算qkv
        self.qkv = nn.Linear(in_features=dim, out_features=dim * 3, bias=qkv_bias)
        # dropout层
        self.atten_drop = nn.Dropout(atten_drop_ratio)
        # 再qkv计算完之后通过一个全连接提取特征
        self.proj = nn.Linear(in_features=dim, out_features=dim)
        # dropout层
        self.proj_drop = nn.Dropout(proj_drop_ratio)

    # 前向传播
    def forward(self, inputs):
        # 获取输入图像的shape=[b,197,768]
        B, N, C = inputs.shape
        # 将输入特征图经过全连接层生成qkv [b,197,768]==>[b,197,768*3]
        qkv = self.qkv(inputs)
        # 维度调整 [b,197,768*3]==>[b, 197, 3, 12, 768//12]
        qkv = qkv.reshape(B, N, 3, self.num_heads, C // self.num_heads)
        # 维度重排==> [3, B, 12, 197, 768//12]
        qkv = qkv.permute(2, 0, 3, 1, 4)
        # 切片提取q、k、v的值，单个的shape=[B, 12, 197, 768//12]
        q, k, v = qkv[0], qkv[1], qkv[2]
        # 针对每个head计算 ==> [B, 12, 197, 197]
        atten = (q @ k.transpose(-2, -1)) * self.scale  # @ 代表在多维tensor的最后两个维度矩阵相乘
        # 对计算结果的每一行经过softmax
        atten = atten.softmax(dim=-1)
        # dropout层
        atten = self.atten_drop(atten)
        # softmax后的结果和v加权 ==> [B, 12, 197, 768//12]
        x = atten @ v
        # 通道重排 ==> [B, 197, 12, 768//12]
        x = x.transpose(1, 2)
        # 维度调整 ==> [B, 197, 768]
        x = x.reshape(B, N, C)
        # 通过全连接层融合特征 ==> [B, 197, 768]
        x = self.proj(x)
        # dropout层
        x = self.proj_drop(x)

        return x

        MLP 层

class MLP(nn.Module):
    # 初始化
    def __init__(self, in_features, hidden_features, out_features=None, drop=0.):
        super(MLP, self).__init__()

        # MLP的输出通道数默认等于输入通道数
        out_features = out_features or in_features
        # 第一个全连接层上升通道数
        self.fc1 = nn.Linear(in_features=in_features, out_features=hidden_features)
        # GeLU激活函数
        self.act = nn.GELU()
        # 第二个全连接下降通道数
        self.fc2 = nn.Linear(in_features=hidden_features, out_features=out_features)
        # dropout层
        self.drop = nn.Dropout(drop)

    # 前向传播
    def forward(self, inputs):
        # [b,197,768]==>[b,197,3072]
        x = self.fc1(inputs)
        x = self.act(x)
        x = self.drop(x)

        # [b,197,3072]==>[b,197,768]
        x = self.fc2(x)
        x = self.drop(x)

        return x

Mask RCNN：

骨干网络：ResNet50+FPN

FPN介绍：

深层网络用于提取语义信息，浅层网络用于提取图像表象信息。如人脸，浅层网络提取包含眼睛鼻子嘴形状轮廓信息，深层网络则提取一张人脸。在检测领域，小目标检测一直是重要任务，深层特征以及浅层特征的融合尤为重要。FPN有效解决了这个问题。图d所示

 特征提取后会送入RPN网络（RegionProposal Network- 区域候选网络）

骨干网络resnet50+fpn提取后的特征 经过3*3卷积分支，上面定位层，返回anchor四个坐标的偏置，下面是分类层，返回前景背景概率。

RPN网络的输出会进入ProposalLayer。首先利用rpn_bbox对anchors进行第一次修正，得到ROI并删除其中的一部分超界的ROI。然后，对剩下的ROI进行score排序，保留其中预测为前景色概率大的一部分（具体值可以在配置文件中进行配置）。最后，利用NMS获得最终的RP。