计图开源: 图匹配网络SuperGlue的优化!速度超TensorRT、显存省一半以上
Part 1
SuperGlue是一个图像特征匹配网络,其输入为两张图片的特征点和描述子,输出为两张图片间的特征匹配关系。其核心技术点是采用注意力机制对特征进行聚合,使其可以推理获取图片中深层全局/上下文信息,并优化特征匹配数量和正确性。
SuperGlue算法在真实室内室外场景的位姿估计任务上可以达到出色的结果,其网络架构主要包含两个部分:注意力图神经网络模块和最优传输匹配模块。
图神经网络模块:借鉴自Transformer的简单变形应用,由特征点位置(position)和描述子(descriptor)共同组成输入嵌入模块(Input Embedding),并在累叠的Self-Attention层(同一图像的内部注意力计算)中额外扩展了Cross-Attention层(不同图像的跨图注意力计算)。在默认全连接网络图注意学习中,该模块的复杂度为O(N^2),在运行时间和显存空间占用上都比通常的卷积上都高出一个数量级。根据其计算量决定性的资源消耗性质,我们对矩阵运算进行分块优化,在不影响并行数的前提下极大降低显存占用。
最优传输匹配模块:如图2所示,原版采用的是Sinkhorn来迭代求解近似的匹配矩阵。我们根据实际实验效果将其替换成更加简化的dual-Softmax算子[2]来优化迭代计算。
图2 Self/Cross注意力机制以及最优传输网络示意(来自SuperGlue原文)
Part 2
该项目通过Jittor开源,开源链接在:
test_gpu/superglue.py:基于计图实现的superglue匹配网络模型,实现了superglue的图片特征匹配算法。 test_gpu/test_time.py:实现计图框架下使用不同优化方法,测试superglue模型进行图片特征匹配的时间性能。 pipeline_inloc_aslfeat.py:将计图版本的superglue整合进视觉定位pipeline,实现了对Inloc数据集中图片的特征匹配,并将匹配结果可视化(见图3)。
# 自动类型转换with jt.flag_scope(auto_mixed_precision_level=5):... 手工类型转换model.float16()
第0级:默认使用float32 第1级:默认使用float32,但开启tensorcore 第2级:同第1级,但前向优化更加激进 第3级:同第2级,前向反向优化都更加激进 第4级:全局使用float16, 但是部分算子使用float32(如求和等) 第5级:同第4级,并且参数也会使用float16 第6级:全局所有算子强制使用float16
欢迎大家使用计图,同时也欢迎加入我们社区进行交流讨论,欢迎为我们提出建议,或者点亮star.
提出issue:https://github.com/Jittor/jittor/issues
Jittor QQ 群:761222083
计图论坛:https://discuss.jittor.org/
参考文献
Sarlin P E, DeTone D, Malisiewicz T, et al,Superglue: Learning feature matching with graph neural networks,Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, 2020, 4938-4947.
Rocco I, Cimpoi M, Arandjelović R, et al, Neighbourhood consensus networks, Advances in neural information processing systems, 2018, 31.
1. CVMJ刊发视觉注意力机制综述,Github 仓库获star1400+
2. Jittor团队提出大核注意力,开源新的骨干网络VAN
4. 首届JDC召开,计图技术委员会评出2021年度优秀Jittor开源项目