“他山之石，可以攻玉”，站在巨人的肩膀才能看得更高，走得更远。在科研的道路上，更需借助东风才能更快前行。为此，我们特别搜集整理了一些实用的代码链接，数据集，软件，编程技巧等，开辟“他山之石”专栏，助你乘风破浪，一路奋勇向前，敬请关注。
作者：知乎—Kissrabbit
地址：https://www.zhihu.com/people/yang-jian-hua-63-91
新的一年，灵机一动（实际是看了NanoDet的工作受了点启发），想再一次尝试YOLO系列的轻量化工作。这次就叫YOLO-Nano吧（嘿？已经有这个名字的工作了啊？哎呀呀，抱歉抱歉，还请原谅~）
这次番外的工作是受了NanoDet工作的启发：
YOLO之外的另一选择，手机端97FPS的Anchor-Free目标检测模型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，模型结构就说完了。再放上来模型代码（只展示必要的部份）：
代码语言：javascript复制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训练我跳票了，实在卡不够用，只能放弃~
本文目的在于学术交流，并不代表本公众号赞同其观点或对其内容真实性负责，版权归原作者所有，如有侵权请告知删除。
