--neozng1@hnu.edu.cn
笔者已经为nanodet增加了非常详细的注释,代码请戳此仓库:nanodet_detail_notes: detail every detail about nanodet 。5. Head 【目标检测算法讲解与部署|NanoDet代码逐行精读与修改(五.1)检测头的构造和前向传播】head部分总共有五百多行代码,将分为
此仓库会跟着文章推送的节奏持续更新!
- 初始化、构造、前向传播
- priors获取、标签分配、loss计算
- 检测框转换、后处理
5.1. 初始化、层构造和前向传播
5.1.1. 参数初始化
class NanoDetPlusHead(nn.Module):
"""Detection head used in NanoDet-Plus.
Args:
num_classes (int): Number of categories excluding the background
category.不包括背景类,可以认为全部类的输出低于某个阈值视为背景
loss (dict): Loss config.
input_channel (int): Number of channels of the input feature.
刚送入检测头的通道数量,需要和PAN的输出通道数量保持一致
feat_channels (int): Number of channels of the feature. 经过检测头卷积后的通道数
Default: 96.
stacked_convs (int): Number of conv layers in the stacked convs. 堆叠的卷积层数
Default: 2.
kernel_size (int): Size of the convolving kernel. Default: 5.
strides (list[int]): Strides of input multi-level feature maps. 下采样步长
Default: [8, 16, 32].
conv_type (str): Type of the convolution.
Default: "DWConv".
norm_cfg (dict): Dictionary to construct and config norm layer.
Default: dict(type='BN').
AGM中使用GN,但是GN对于CPU型设备不友好,BN可以直接和卷积一同参数化而不需要额外计算
reg_max (int): The maximal value of the discrete set. Default: 7.
参见GFL的论文,用于建模框的任意分布,而不是得到一个Dirac分布
在第四部分的AGM中也介绍过,请查看往期博客
activation (str): Type of activation function. Default: "LeakyReLU".
assigner_cfg (dict): Config dict of the assigner. Default: dict(topk=13).
"""
def __init__(
self,
num_classes,
loss,
input_channel,
feat_channels=96,
stacked_convs=2,
# 选择5x5的卷积大核
kernel_size=5,
# 输入特征相对原图像的下采样率,因为PAN中采用了extra_layer-
# 从配置文件也可以看到这里实际有四层:[8,16,32,64]
strides=[8, 16, 32],
conv_type="DWConv",
norm_cfg=dict(type="BN"),
reg_max=7,
activation="LeakyReLU",
assigner_cfg=dict(topk=13),
**kwargs
):
super(NanoDetPlusHead, self).__init__()
self.num_classes = num_classes
self.in_channels = input_channel
self.feat_channels = feat_channels
self.stacked_convs = stacked_convs
self.kernel_size = kernel_size
self.strides = strides
self.reg_max = reg_max
self.activation = activation
# 必然是使用深度可分离卷积了,DepthwiseConvModule来自MMDetection,请直接看源码
self.ConvModule = ConvModule if conv_type == "Conv" else DepthwiseConvModule
?
self.loss_cfg = loss
self.norm_cfg = norm_cfg
?
# 第四部分介绍的动态分配器,稍后计算loss会使用到
self.assigner = DynamicSoftLabelAssigner(**assigner_cfg)
# 根据输出的框分布进行积分,得到最终的位置值
self.distribution_project = Integral(self.reg_max)
?
# 联合了分类和框的质量估计表示
self.loss_qfl = QualityFocalLoss(
beta=self.loss_cfg.loss_qfl.beta,
loss_weight=self.loss_cfg.loss_qfl.loss_weight,
)
# 初始化参数中reg_max的由来,在对应模块中进行了详细的介绍
self.loss_dfl = DistributionFocalLoss(
loss_weight=self.loss_cfg.loss_dfl.loss_weight
)
# IoU loss的一种改进,IoU loss家族还有CIoU/DIoU等
self.loss_bbox = GIoULoss(loss_weight=self.loss_cfg.loss_bbox.loss_weight)
self._init_layers()
self.init_weights()
5.1.2. 构造和权重设置
下面是
_buid_not_shared_head
、_init_layers()
和init_weights()
:def _buid_not_shared_head(self):
cls_convs = nn.ModuleList()
# stacked_convs是参数中设定的卷积层数
for i in range(self.stacked_convs):
# 第一层要和PAN的输出对齐通道
chn = self.in_channels if i == 0 else self.feat_channels
cls_convs.append(
self.ConvModule(
chn,
self.feat_channels,
self.kernel_size,
stride=1,
# 加大小为卷积核一般的padding使得输入输出feat有相同尺寸
padding=self.kernel_size // 2,
norm_cfg=self.norm_cfg,
bias=self.norm_cfg is None,
activation=self.activation,
)
)
return cls_convs
def _init_layers(self):
self.cls_convs = nn.ModuleList()
for _ in self.strides:
# 为每个stride的创建一个head,cls和reg共享这些参数
cls_convs = self._buid_not_shared_head()
self.cls_convs.append(cls_convs)
?
# 同样,为每个头增加gfl卷积
self.gfl_cls = nn.ModuleList(
[
nn.Conv2d(
self.feat_channels,
# 每个位置需要num_classes个通道用于预测类别分数,还有4*(reg_max+1)来回归位置
# 用同一组卷积来获得,输出结果时再split成两份即可
self.num_classes + 4 * (self.reg_max + 1),
1,
padding=0,
)
for _ in self.strides
]
)
还有平平无奇的权重初始化:
# 采用norm初始化
def init_weights(self):
for m in self.cls_convs.modules():
if isinstance(m, nn.Conv2d):
normal_init(m, std=0.01)
# init cls head with confidence = 0.01
bias_cls = -4.595
for i in range(len(self.strides)):
normal_init(self.gfl_cls[i], std=0.01, bias=bias_cls)
print("Finish initialize NanoDet-Plus Head.")
5.1.3. 前向传播
# head的推理方法
def forward(self, feats):
# 有一个为了兼容onnx的方法
if torch.onnx.is_in_onnx_export():
return self._forward_onnx(feats)
# 输出默认有4份,是一个list
outputs = []
# feats来自fpn,有多组,且组数和self.cls_cons/self.gfl_cls的数量需保持一致
# 默认的参数设置是4组
for feat, cls_convs, gfl_cls in zip(
feats,
self.cls_convs,
self.gfl_cls,
):
# 对每组feat进行前向推理操作
for conv in cls_convs:
feat = conv(feat)
output = gfl_cls(feat)
# 所有head的输出会在展平后拼接成一个tensor,方便后处理
# output是一个四维tensor,第一维长度为1(=batch size)
# 长为W宽为H(其实长宽相等)即feat的大小,高为80 + 4 * (reg_max+1)即cls和reg
# 按照第三个维度展平,就是排成一个长度为W*H的tensor,另一个维度是输出的cls和reg
outputs.append(output.flatten(start_dim=2)) # 变成1x112x(W*H)维了,80+4*8=112
# 把不同head的输出交换一下维度排列顺序,全部拼在一起
# 按照第三维拼接,就是1x112x2125(对于nanodet-m)
outputs = torch.cat(outputs, dim=2).permute(0, 2, 1)
return outputs
因为第一个维度长度为1的维度没有消除掉,所以很多读者初看认为这是一个三维的向量,这造成了一些困惑。对于训练或者批量推理的时候,第一个向量的长度就不会为1了!我们直接看看NanoDet-m转成onnx后的可视化:
文章图片
上图是四个头的输出
?split操作不需要理会,这是在推理的时候拆分分类输出和位置输出使用到的算子。我们会在介绍完训练后介绍部署和推理。
推荐阅读
- 语义分割|cv2 interpolate插值-align_corners
- 大数据|升级版NanoDet-Plus来了!简单辅助模块加速训练收敛,精度大幅提升!
- 目标检测算法讲解与部署|NanoDet代码逐行精读与修改(四)动态软标签分配(dynamic soft label assigner)
- 3D|点云边界提取方法总结
- CV|2. vit——vision transformer
- 深度学习|yolov5安全帽检测、反光衣检测、抽烟检测
- 深度学习|关于添加注意力模块的yolov5-5.0与yolov5-6.0的检测效果对比
- python|python批量修改txt文件里的类别数,批量修改文件名
- 深度学习|yolov5-v6.1发布