用更小更快的模型达到state-of-the-art效果的多任务学习
今天跟大家分享几天前arXiv上的一篇兼顾模型速度与性能的论文《Real-Time Joint Semantic Segmentation and Depth Estimation Using Asymmetric Annotations》,来自阿德莱德大学和墨尔本大学的研究人员在该文解决了在计算受限的设备上部署多任务模型的问题,取得了又好(达到state-of-the-art效果)又快(17毫秒每帧)同时实现语义分割与深度估计的效果。 作者信息:
在机器人上部署深度学习模型提取环境信息是一项艰巨的任务,即使是使用强大的通用的GPU。在这篇论文中,作者希望将语义分割与深度估计两个感知任务部署到机器人中,其主要解决了三个问题: 1)单模型多任务,在同一个深度学习网络中同时实现语义分割与深度估计; 2)实时计算; 3)非对称的标注数据,即并不是所有标注数据既有语义标注又有深度标注。
网络架构与算法设计思想
1)骨干网选择和改进: 作者使用构建于MobileNet-v2分类网络上的Light-Weight RefineNet(BMVC2018论文提出的)作为骨干网,该网络在输入图像大小为640 × 480时计算量14 GFLOPs,为进一步降低计算量,作者将最后的CRP block中的1*1卷积替换为分组卷积(grouped convolution),使其降低为6.5 GFLOPs。 2)语义分割与深度估计联合学习: 在上述Light-Weight RefineNet结构之后,网络分成两个预测任务,分别使1*1的depthwise卷积和3*3的普通的卷积。 其损失函数:
3)知识蒸馏用于非对称标注数据的专家标注: 这一步是重点!对于多任务模型,需要标注的数据含有每种任务的标签才能训练。对于两种任务T1和T2,假设只有少部分数据被标注了两种标签,对于样本数量更多的只有其中一类标签的数据,引入一个更加强大的专家模型,计算其在另一任务中的预测结果作为合成ground truth数据。 训练的时候,使用合成ground truth数据先预训练网络,然后再使用拥有两种真实标签的ground truth对网络fine-tune。
实验结果 作者使用的机器配置和软件版本24GB RAM,Intel i5-7600 处理器,一块GT1080Ti GPU,CUDA9.0,CuDNN7.0。 1)NYUDv2数据集实验 作者首先在NYUDv2数据集上做了实验,其含有1449图像同时含有语义和深度标注(795幅训练集,654幅验证集),同时又有超过300000幅图像有深度标注。 作者引入的专家模型是Light-Weight RefineNet-152 语义分割模型,其在验证集上的精度是44.4% mean iou。 制作完合成语义标注后,使用大的含有合成标注的数据集预训练,然后在795图像的小数据集上fine-tune。 结果如下:
参数量仅3.07M,打败了多个大型网络模型,预测性能几乎都接近最好的结果,计算时间仅需要12.8ms。 下图是深度估计的精度进一步比较:
接近大型网络的结果,甚至在某些指标下超越了大型网络的结果。 预测结果示例:
作者同时比较了在小数据集上的单任务模型,和其他文献中使用合成标签训练的预测结果,表明该文方法取得了明显的优势。
2)KITTI数据集实验 KITTI数据集仅含有146幅图像有语义标注(100幅训练46幅测试),同时含有20697幅图像有深度标注(20000幅训练,697幅测试),他们之间没有同时被标注了语义和深度的数据。 作者使用在CityScapes数据集上训练的ResNet-38模型作为语义专家模型给20000仅有深度标注的图像预测语义标签,使用本文提出的网络架构在20000幅具有深度标注的图像上训练的深度估计模型,然后给100幅仅有语义标注的图像预测深度标签。 按照上面的方法,现在大库上预训练,再在100幅图像的小库上fine-tune,结果如下:
比较有意思的是,这两个任务甚至比目前state-of-the-art的大型网络都达到了更好的结果。参数量仅2.99M,而且在1200*350的图像上一帧计算时间仅需要16.9毫秒。 更细节的分类比较:
预测示例:
3)更多扩展
作者又研究了扩展到单网络三种任务(Single Model – Three Tasks)、单网络两个数据集两种任务(Single Model – Two Datasets, Two Tasks)和集成到Dense Semantic SLAM,都取得了又快又好地结果。
很可惜,该文目前并无开源代码。
论文地址: https://arxiv.org/abs/1809.04766