YOLO9000:更好、更快、更强的实时目标检测系统
一、研究背景与目标
本文由Joseph Redmon与Ali Farhadi共同撰写,发表于2017年IEEE计算机视觉与模式识别会议(CVPR),是目标检测领域具有里程碑意义的经典论文之一。作者分别来自华盛顿大学、艾伦人工智能研究所(Allen Institute for AI)以及XNOR.AI。该项研究的核心目标是在保持实时检测速度的前提下,大幅提升目标检测系统的准确率与可识别类别数量,最终构建出一个能够检测超过9000类物体类别的系统。
在当时的学术背景下,目标检测系统主要面临两个突出问题:其一是检测系统的类别覆盖范围极为有限。尽管ImageNet等分类数据集已经包含数百万张图像和成千上万个类别,但主流检测数据集(如PASCAL VOC和COCO)通常仅包含几十到数百个类别的边界框标注。检测标注远比分类标注昂贵,因此短期内难以出现大规模检测数据集。其二是检测系统的速度与精度难以兼顾。基于区域提议(region proposal)的方法虽然精度较高,但计算开销巨大,难以实现实时处理。YOLO(You Only Look Once)虽然速度快,但在定位精度和召回率上仍有不足。
基于这一背景,作者提出了两项主要贡献:第一,通过一系列改进将原始YOLO方法升级为YOLOv2,显著提升其检测性能;第二,提出一种联合训练检测与分类数据的方法,并借助WordTree层级结构将不同数据集的标签体系融合起来,从而训练出能够检测超过9000个类别的YOLO9000系统。
二、YOLOv2的主要改进
为了使检测器更准确且仍然保持高速,作者从多个角度对YOLO进行了系统改进。论文详细列举了每一项改进策略及其带来的性能提升效果。
首先,作者引入了批归一化(batch normalization)。在所有卷积层后加入批归一化后,模型在PASCAL VOC 2007上的平均精度(mAP)提升了超过2%,同时可以移除dropout而不产生过拟合。批归一化通过稳定各层输入的分布,加速了收敛并起到正则化作用。
第二项改进是高分辨率分类器(high resolution classifier)的微调策略。原始YOLO在224×224分辨率上预训练分类网络,再直接切换至448×448进行检测训练,导致网络需要同时适应分辨率变化和检测任务。YOLOv2先在448×448分辨率下对分类网络微调10个epoch,再进行检测训练,这一策略使mAP提升了近4%。
第三项重要改进是引入锚框(anchor boxes)机制。Faster R-CNN使用人工预设的锚框来预测边界框偏移量而非直接预测坐标,这简化了学习问题。YOLOv2移除了全连接层,改用锚框预测边界框,并将网络输入从448×448调整为416×416,以获得13×13的奇数特征图,使大物体能够由中心单元负责预测。尽管mAP从69.5微降至69.2,但召回率(recall)从81%提升到88%,为后续改进提供了更大空间。
第四项改进是维度聚类(dimension clusters)。传统锚框的尺寸是人工设定的,作者改用K-means聚类算法对训练集中的边界框尺寸进行聚类,以自动获得更好的先验框(priors)。为了避免大框主导欧氏距离,聚类时采用了基于交并比(IoU)的距离度量:d(box, centroid) = 1 − IoU(box, centroid)。实验表明,仅使用5个聚类中心即可达到与9个人工锚框相当的平均IoU(61.0 vs 60.9),而使用9个聚类中心时平均IoU提升至67.2%。
第五项改进是直接位置预测(direct location prediction)。原始锚框方法中,网络预测的无约束偏移量可能导致锚框在图像中任意位置出现,造成训练初期不稳定。YOLOv2改为预测相对于网格单元的位置坐标,并通过sigmoid函数将中心坐标约束在0到1之间。边界框的宽高则通过聚类先验乘以指数化预测值获得。该方法使网络更稳定,并在维度聚类基础上将mAP提升了近5%。
第六项改进是细粒度特征(fine-grained features)的引入。为了改善对较小物体的检测能力,作者增加了一个直通层(passthrough layer),将26×26分辨率的浅层特征图重组为13×13×2048并与深层特征拼接,使检测器能够同时利用高分辨率与低分辨率的特征信息。这一改进带来了约1%的性能提升。
第七项改进是多尺度训练(multi-scale training)。由于YOLOv2仅包含卷积层和池化层,网络可以动态调整输入尺寸。作者在训练过程中每10个批次随机选择320×320到608×608之间的不同输入尺寸(均为32的倍数),使网络学会在不同分辨率下进行预测。这使得同一个模型能够在运行时根据实际需求在速度与精度之间灵活权衡。在低分辨率288×288下,系统能以超过90 fps运行,而在544×544下则可达到78.6 mAP,同时仍保持40 fps的实时速度。
三、Darknet-19网络架构
为了进一步提升速度,作者设计了新的基础网络Darknet-19。该网络包含19个卷积层和5个最大池化层。其设计融合了VGG模型常用的3×3卷积核与逐池化后通道数翻倍的策略,同时借鉴了Network in Network的思想,使用1×1卷积压缩特征表示,并采用全局平均池化。Darknet-19在ImageNet上仅需55.8亿次浮点运算即可取得72.9%的top-1准确率和91.2%的top-5准确率,远优于需要306.9亿次运算的VGG-16。在检测任务上,作者移除了最后一个卷积层并加入三个1024滤波器的3×3卷积层和一个1×1卷积层用于输出检测预测。
四、WordTree与联合训练
为了解决检测数据集类别少而分类数据集类别多的问题,作者提出了基于WordNet构建的层级分类结构——WordTree。WordNet是一个有向图结构的知识库,其中“Norfolk terrier”和“Yorkshire terrier”都是“terrier”的下位词,“terrier”又属于“hunting dog”,再属于“dog”,最终归至“physical object”。由于WordNet是图而非树,作者选择最短路径原则将其简化为层级树。使用WordTree进行分类时,网络在每个节点预测该节点各个下位词的条件概率,绝对概率由根节点到目标节点的路径上所有条件概率相乘得到。例如,要判断图像是否为Norfolk terrier,需要计算Pr(Norfolk terrier|terrier)×Pr(terrier|hunting dog)×…×Pr(animal|physical object)。在ImageNet 1000类基础上加入所有中间节点后,WordTree1k的标签空间扩展至1369个类别,但top-1准确率仅从72.9%轻微下降至71.9%,显示出该方法对标签结构扩展具有很强的鲁棒性。
利用WordTree,作者将COCO检测数据集与ImageNet分类数据集的前9000个类别融合为一个包含9418个节点的层级结构。在联合训练过程中,当网络看到检测图像时,使用完整的YOLOv2损失函数进行反向传播;当看到分类图像时,仅反向传播分类相关部分的损失。对于分类图像,系统找到对该类别预测概率最高的边界框,并仅在该框上计算分类损失和基于IoU≥0.3假设的目标性损失。
五、实验结果与系统性能
在PASCAL VOC 2007测试集上,YOLOv2在416×416输入下取得76.8 mAP,速度达到67 fps;在544×544输入下取得78.6 mAP,速度仍保持在40 fps,超过了Faster R-CNN with ResNet和SSD的精度,同时速度远快于这些方法。在PASCAL VOC 2012测试集上,YOLOv2取得73.4 mAP,与当时最先进的检测器性能相当,但速度快2至10倍。在COCO test-dev2015上,YOLOv2在VOC指标(IoU=0.5)下取得44.0 mAP。
YOLO9000在ImageNet检测验证集上取得了19.7 mAP的总体成绩,其中156个从未见过检测标注的类别上取得16.0 mAP。该成绩超过了DPM方法,且YOLO9000同时还能实时检测超过9000个类别。从类别表现来看,YOLO9000对动物类新物种学习效果较好,例如对tiger的AP达到61.0,对koala bear达到54.3;但在服装和装备类上表现不佳,如sunglasses和swimming trunks的AP均为0.0。这是因为COCO中缺乏服装类的边界框标注,导致目标性预测难以泛化至这些类别。
六、总结与价值
本文通过系统性的改进构建了YOLOv2,使其成为当时最先进的实时目标检测器之一。同时,作者首次提出了利用WordTree将检测数据与分类数据联合训练的方法,成功训练出能够检测超过9000个类别的YOLO9000系统。该研究显著缩小了检测与分类在类别覆盖范围上的差距,为弱监督目标检测和大规模视觉识别提供了新的思路。WordTree的层级表示方法不仅适用于检测任务,也有望推广到图像分类和语义分割等领域。