--neozng1@hnu.edu.cn

NanoDet是一个单阶段的anchor-free模型,其设计基于FCOS模型,并加入了动态标签分配策略/GFL loss和辅助训练模块。由于其轻量化的设计和非常小的参数量,在边缘设备和CPU设备上拥有可观的推理速度。其代码可读性强扩展性高,是目标检测实践进阶到深入的不二选择。作者在知乎上有一篇介绍的文章,指路:超简单辅助模块加速训练收敛,精度大幅提升!移动端实时的NanoDet升级版NanoDet-Plus来了!
笔者已经为nanodet增加了非常详细的注释,代码请戳此仓库:
此仓库会跟着文章推送的节奏持续更新!

话不多说,从结构上我们直接分backbone、neck、head、assist module、dynamic label assigner五个模块进行超級详细的介绍。

0. Architecture

surprise!首先当然要介绍一下整体的架构了。先看看整个模型的架构图:

NanoDet-Plus架构图,图源作者知乎

直观来看,最大的不同就是新增的Assign Guidance Module模块。检测框架还是FCOS式的一阶段网络,neck改为了GhostPAN,同时摒弃了FCOS的标签分配策略转向动态软标签分配并加入辅助训练模块也就是前述的AGM,它将作为教师模型帮助head获得更好的训练。头部的回归和标签预测仍然继承之前的Generalized Focal Loss。

以NanoDet-m (320x320)为例让我们先看一下config/下的配置文件中和网络架构有关的选项:

name: NanoDetPlus
    detach_epoch: 10
    backbone:
      name: ShuffleNetV2    # 默认使用shuffleNetV2
      model_size: 1.0x      # 模型缩放系数,更大的模型就是相应扩大各层feature map的大小
      out_stages: [2,3,4]   # backbone中输出特征到FPN的stage
      activation: LeakyReLU # 激活函数采用LeakyReLU
    fpn:
      name: GhostPAN        # 用ghostNet的模块对不同特征层进行融合
      in_channels: [116, 232, 464]  # 输入fpn的feature map尺寸
      out_channels: 96      # 
      kernel_size: 5        # 卷积核大小
      num_extra_level: 1    # 输出额外一层,即在最底层的基础上再上采样得到更大的feature map
      use_depthwise: True   # 是否使用深度可分离卷积
      activation: LeakyReLU # 激活函数
    head:
      name: NanoDetPlusHead # 检测头,还提供了之前的nanodet头和最简单的卷积头
      num_classes: 80	    # 类别数
      input_channel: 96     # 输入通道数
      feat_channels: 96     # 和输入通道数一致
      stacked_convs: 2		# 头的卷积层数
      kernel_size: 5   		# nanodet-plus也换用了5x5的大核卷积
      strides: [8, 16, 32, 64] # 有四个头,分别对应不同尺度特征的检测,这是不同head检测时的下采样倍数
      activation: LeakyReLU 
      reg_max: 7  # 用于DFL的参数,head的回归分支会预测框的分布,即用回归reg_max+1个离散的几个值来表示一个分布            
      norm_cfg: 
        type: BN	# head选用Batch Norm进行归一化操作
      loss:
      	# loss仍然继承了nanodet,使用GFL,并且这些loss有不同的权重
        loss_qfl:   
          name: QualityFocalLoss
          use_sigmoid: True
          beta: 2.0
          loss_weight: 1.0
        loss_dfl:
          name: DistributionFocalLoss
          loss_weight: 0.25
        loss_bbox:
          name: GIoULoss   # 选取计算IOU loss的方法为GIoU
          loss_weight: 2.0
    # Auxiliary head, only use in training time.
    # 新增的辅助训练模块,其实就是一个常规的检测头(上面的头是简化过的版本,表达能力显然不如标准头)
    aux_head:
      name: SimpleConvHead
      num_classes: 80
      input_channel: 192 # 可以看到输入通道数也比nanodet head多
      feat_channels: 192
      stacked_convs: 4   # 堆叠4层而不是上面的2层,反正是训练不是推理
      strides: [8, 16, 32, 64] # 对应nanodet head的四个头
      activation: LeakyReLU
      reg_max: 7  # 同head中的参数

下图是训练时feature的流图,backbone的输出进入两个Ghost PAN,其中一个是为AGM专门搭建的,另一个pan和head连接。AGM会将两个PAN的输出拼接在一起作为输入(猜想这样可以更好的获取当前Head的训练情况,同时也能获取更多特征),其有两个分支,分别负责生成用作标签分配的cls_pred和reg_pred。对于Ghost PAN中的不同层特征,AGM采用相同的参数(参数共享)进行运算,大大减小了训练时的参数数和运算量,提升精度的同时使得训练速度不会增加太多。AGM的输出在训练初期将会作为Head标签分配的参考,并且AGM的loss也会进行回传,帮助网络更快地收敛。经过数个epoch(默认是10个)的训练后Head的预测已经有较好的准确度,此时将AGM模块分离,直接由Head的输出自行完成标签分配的任务。

在训练完成进行推理时,直接去除AGM和aux_fpn,得到非常精简的网络结构。