【他山之石】超轻量的YOLO-Nano

“他山之石，可以攻玉”，站在巨人的肩膀才能看得更高，走得更远。在科研的道路上，更需借助东风才能更快前行。为此，我们特别搜集整理了一些实用的代码链接，数据集，软件，编程技巧等，开辟“他山之石”专栏，助你乘风破浪，一路奋勇向前，敬请关注。

作者：知乎—Kissrabbit

地址：https://www.zhihu.com/people/yang-jian-hua-63-91

新的一年，灵机一动（实际是看了NanoDet的工作受了点启发），想再一次尝试YOLO系列的轻量化工作。这次就叫YOLO-Nano吧（嘿？已经有这个名字的工作了啊？哎呀呀，抱歉抱歉，还请原谅~）

这次番外的工作是受了NanoDet工作的启发：

依据NanoDet模型，我将我的YOLOv3Slim的backbone换成了十分轻量的ShuffleNetv2，并且去掉最后一层1024的卷积层，而head部份的设计和NanoDet是一样的，十分轻量的PAN结构。至于loss设计，目前仍用YOLO系列的，暂未改动。

01

模型结构

模型结构图：

YOLO-Nano网络结构

简单说一下：

1.1 Backbone

使用的是ShuffleNetv2，相关代码我借鉴了NanoDet所提供的网络代码和模型下载地址，由于ShuffleNetv2有大量的depthwise卷积，因此如果不做特殊处理的话，可能对GPU不是太友好。但这样的模型本身是面向嵌入式等移动平台的，鲜有诸如1080ti、2080ti这样的gpu算力。

关于backbone代码，大家可以打开我项目的backbone/shufflenetv2.py。主要使用shufflenetv2-0.5x和shufflenetv2-1.0x两个模型。

1.2 neck

拟定使用spp，这个还没有决定是否加进去，因此待定中。悄悄说一句，应该没有人想尝试家DCNv2吧~那玩意挺慢的。

另外，还用了PAN，PAN的设计参考了NanoDet：

a、去掉所有的卷积，仅仅保留FPN中必要的1x1卷积进行通道对齐，这里，我们将三个尺度的feature map的通道都用1x1卷积处理成96。为什么是96，可以参考NanoDet的设计。

b、所有的上采样操作（对应FPN）、下采样操作（对应PAN后续部份）均使用插值：

# FPN
p4 = self.smooth_0(p4 + F.interpolate(p5, scale_factor=2.0))
p3 = self.smooth_1(p3 + F.interpolate(p4, scale_factor=2.0))

# PAN
p4 = self.smooth_2(p4 + F.interpolate(p3, scale_factor=0.5))
p5 = self.smooth_3(p5 + F.interpolate(p4, scale_factor=0.5))

c、特征融合使用sum操作，而不是cat操作。

d、特征融合后，额外接一个3x3卷积再处理一下，这一是与NanoDet不同的。按照NanoDet的做法，特征融合几乎全是线性操作，个人觉得不太妥，因此，在考虑了速度和精度后，决定加入一层3x3的卷积（后接BN+LeakyReLU）来处理一下，速度上带来的额外开销很小，精度提升了2-3mAP（VOC test）。

1.3 head

head的设计和NanoDet一样：3x3的dw卷积+1x1普通卷积，一共重复两次，如上图所示，没有复杂的东西。

ok，模型结构就说完了。再放上来模型代码（只展示必要的部份）：

class YOLONano(nn.Module):
    def __init__(self, device, input_size=None, num_classes=20, trainable=False, conf_thresh=0.001, nms_thresh=0.50, anchor_size=None, backbone='1.0x', diou_nms=False):
        super(YOLONano, self).__init__()
        self.device = device
        self.input_size = input_size
        self.num_classes = num_classes
        self.trainable = trainable
        self.conf_thresh = conf_thresh
        self.nms_thresh = nms_thresh
        self.nms_processor = self.diou_nms if diou_nms else self.nms
        self.bk = backbone
        self.stride = [8, 16, 32]
        self.anchor_size = torch.tensor(anchor_size).view(3, len(anchor_size) // 3, 2)
        self.anchor_number = self.anchor_size.size(1)

        self.grid_cell, self.stride_tensor, self.all_anchors_wh = self.create_grid(input_size)
        self.scale = np.array([[[input_size[1], input_size[0], input_size[1], input_size[0]]]])
        self.scale_torch = torch.tensor(self.scale.copy(), device=device).float()

        if self.bk == '0.5x':
            # use shufflenetv2_0.5x as backbone
            print('Use backbone: shufflenetv2_0.5x')
            self.backbone = shufflenetv2(model_size=self.bk, pretrained=trainable)
            width = 0.4138
        elif self.bk == '1.0x':
            # use shufflenetv2_1.0x as backbone
            print('Use backbone: shufflenetv2_1.0x')
            self.backbone = shufflenetv2(model_size=self.bk, pretrained=trainable)
            width = 1.0
        else:
            print("For YOLO-Nano, we only support <0.5x, 1.0x> as our backbone !!")
            exit(0)


         # FPN+PAN
        self.conv1x1_0 = Conv(int(116*width), 96, k=1)
        self.conv1x1_1 = Conv(int(232*width), 96, k=1)
        self.conv1x1_2 = Conv(int(464*width), 96, k=1)

        self.smooth_0 = Conv(96, 96, k=3, p=1)
        self.smooth_1 = Conv(96, 96, k=3, p=1)
        self.smooth_2 = Conv(96, 96, k=3, p=1)
        self.smooth_3 = Conv(96, 96, k=3, p=1)

        # det head
        self.head_det_1 = nn.Sequential(
            Conv(96, 96, k=3, p=1, g=96),
            Conv(96, 96, k=1),
            Conv(96, 96, k=3, p=1, g=96),
            Conv(96, 96, k=1),
            nn.Conv2d(96, self.anchor_number * (1 + self.num_classes + 4), 1)
        )
        self.head_det_2 = nn.Sequential(
            Conv(96, 96, k=3, p=1, g=96),
            Conv(96, 96, k=1),
            Conv(96, 96, k=3, p=1, g=96),
            Conv(96, 96, k=1),
            nn.Conv2d(96, self.anchor_number * (1 + self.num_classes + 4), 1)
        )
        self.head_det_3 = nn.Sequential(
            Conv(96, 96, k=3, p=1, g=96),
            Conv(96, 96, k=1),
            Conv(96, 96, k=3, p=1, g=96),
            Conv(96, 96, k=1),
            nn.Conv2d(96, self.anchor_number * (1 + self.num_classes + 4), 1)
        )


    def forward(self, x, target=None):
        # backbone
        c3, c4, c5 = self.backbone(x)

        # # neck
        # c5 = self.spp(c5)

        p3 = self.conv1x1_0(c3)
        p4 = self.conv1x1_1(c4)
        p5 = self.conv1x1_2(c5)

        # FPN
        p4 = self.smooth_0(p4 + F.interpolate(p5, scale_factor=2.0))
        p3 = self.smooth_1(p3 + F.interpolate(p4, scale_factor=2.0))

        # PAN
        p4 = self.smooth_2(p4 + F.interpolate(p3, scale_factor=0.5))
        p5 = self.smooth_3(p5 + F.interpolate(p4, scale_factor=0.5))


        # det head
        pred_s = self.head_det_1(p3)
        pred_m = self.head_det_2(p4)
        pred_l = self.head_det_3(p5)

        preds = [pred_s, pred_m, pred_l]
        total_conf_pred = []
        total_cls_pred = []
        total_txtytwth_pred = []
        B = HW = 0
        for pred in preds:
            B_, abC_, H_, W_ = pred.size()

            # [B, anchor_n * C, H, W] -> [B, H, W, anchor_n * C] -> [B, H*W, anchor_n*C]
            pred = pred.permute(0, 2, 3, 1).contiguous().view(B_, H_*W_, abC_)

            # Divide prediction to obj_pred, xywh_pred and cls_pred   
            # [B, H*W*anchor_n, 1]
            conf_pred = pred[:, :, :1 * self.anchor_number].contiguous().view(B_, H_*W_*self.anchor_number, 1)
            # [B, H*W*anchor_n, num_cls]
            cls_pred = pred[:, :, 1 * self.anchor_number : (1 + self.num_classes) * self.anchor_number].contiguous().view(B_, H_*W_*self.anchor_number, self.num_classes)
            # [B, H*W*anchor_n, 4]
            txtytwth_pred = pred[:, :, (1 + self.num_classes) * self.anchor_number:].contiguous()

            total_conf_pred.append(conf_pred)
            total_cls_pred.append(cls_pred)
            total_txtytwth_pred.append(txtytwth_pred)
            B = B_
            HW += H_*W_

        conf_pred = torch.cat(total_conf_pred, 1)
        cls_pred = torch.cat(total_cls_pred, 1)
        txtytwth_pred = torch.cat(total_txtytwth_pred, 1).view(B, -1, 4)

        # train
        if self.trainable:
            txtytwth_pred = txtytwth_pred.view(B, HW, self.anchor_number, 4)

            x1y1x2y2_pred = (self.decode_boxes(txtytwth_pred) / self.scale_torch).view(-1, 4)
            x1y1x2y2_gt = target[:, :, 7:].view(-1, 4)

            # compute iou
            iou_pred = tools.iou_score(x1y1x2y2_pred, x1y1x2y2_gt, batch_size=B)

            txtytwth_pred = txtytwth_pred.view(B, -1, 4)
            # compute loss
            conf_loss, cls_loss, bbox_loss, total_loss = tools.loss(pred_conf=conf_pred,
                                                                    pred_cls=cls_pred,
                                                                    pred_txtytwth=txtytwth_pred,
                                                                    label=target,
                                                                    num_classes=self.num_classes
                                                                    )


            return conf_loss, cls_loss, bbox_loss, total_loss

        # test
        else:
            txtytwth_pred = txtytwth_pred.view(B, HW, self.anchor_number, 4)
            with torch.no_grad():
                # batch size = 1                
                all_obj = torch.sigmoid(conf_pred)[0]           # 0 is because that these is only 1 batch.
                all_bbox = torch.clamp((self.decode_boxes(txtytwth_pred) / self.scale_torch)[0], 0., 1.)
                all_class = (torch.softmax(cls_pred[0, :, :], dim=1) * all_obj)
                # separate box pred and class conf
                all_obj = all_obj.to('cpu').numpy()
                all_class = all_class.to('cpu').numpy()
                all_bbox = all_bbox.to('cpu').numpy()

                bboxes, scores, cls_inds = self.postprocess(all_bbox, all_class)

                # print(len(all_boxes))
                return bboxes, scores, cls_inds

02

label制作

大体上同我的yolov3，但有一个小改动，那就是一个gt可能匹配多个anchor box，而不再是只保留iou最大的，目的就是为了增加正样本数量。其他的匹配原则没有变化。后续会尝试使用更好的匹配方式，如ATSS。

03

Loss函数

大体上同我的yolov3，但有一个小改动：obj的目标不再是动态的iou，而是简单的01二分类。这是我亲自试出来的，对于这种小模型，01二分类反而性能要更好一点。其他就都保持一样了，没改动。

另外，我在考虑如何把obj，cls以及iou换个方式耦合起来共同优化

不过，我发现obj的loss是模型的瓶颈，很难下去，这就导致了漏检现象很严重，这一问题我感觉和NanoDet中提到的有关Centerness问题是类似的，所以，如何有效提高obj的收敛速度，对于提升模型性能是有很大帮助的。

04

实验

4.1 VOC

目前voc上已经跑完了：

416下，FPS是13-20fps（i9-9940k），速度测试完全就是拿pytorch-cpu版硬跑~笔者目前还不会诸如nvcc等部署和其他的加速操作。所以这一块没有什么秘密。

模型文件（.pth）只有5M，很小，所以这次就直接传到github了，方便下载。

4.2 COCO

coco：

BFlops：1B

Params：1.33M

还不错，不比tiny-yolo-v3差。

放几张在COCO-val上的可视化结果（当然，我肯定是挑几张好的放出来喽~）

05

项目链接

https://github.com/yjh0410/YOLO-Nano

不是多么出色的项目，所以就不奢求 STAR 了~

06

优化工作

基础工作是完成了，后面再尝试做各种优化，看看能不能把性能再往上拔一拔，实在不行就放弃，掉头发。

后续我要尝试一下nvcc部署，把这模型放到安卓手机上~

PS：之前CSPDarknet的backbone训练我跳票了，实在卡不够用，只能放弃~

本文目的在于学术交流，并不代表本公众号赞同其观点或对其内容真实性负责，版权归原作者所有，如有侵权请告知删除。

“他山之石”历史文章

更多他山之石专栏文章，

请点击文章底部“阅读原文”查看

分享、点赞、在看，给个三连击呗！