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。