计图开源: 图匹配网络SuperGlue的优化！速度超TensorRT、显存省一半以上

2022/04/02 Tutorial

计图开源: 图匹配网络SuperGlue的优化！速度超TensorRT、显存省一半以上

图像特征匹配是SLAM、三维重建以及长时间视觉定位的基础。传统的匹配方式是由局部特征本身的描述子决定的最相似点匹配策略(Nearest Neighbor)，难以应对大视角变化、时间环境变化，以及重复纹理等挑战场景。
SuperGlue[1]是针对大规模特征点的图匹配深度学习网络框架，其利用Transformer的注意力机制来自适应增强特征点的全局信息，提高匹配数和正确性。但同时，注意力机制也引入了更多的计算量，给显存占用和运行实时性带来了严峻的挑战。我们基于Jittor框架，从内存占用和运行时间两个角度对其全面优化。

图1 传统NN匹配和图匹配结果对比

Part 1

SuperGlue简介

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开源，开源链接在：

https://github.com/smy-THU/superglue-jittor
项目代码组织形式如下：
  • test_gpu/superglue.py：基于计图实现的superglue匹配网络模型，实现了superglue的图片特征匹配算法。
  • test_gpu/test_time.py：实现计图框架下使用不同优化方法，测试superglue模型进行图片特征匹配的时间性能。
  • pipeline_inloc_aslfeat.py：将计图版本的superglue整合进视觉定位pipeline，实现了对Inloc数据集中图片的特征匹配，并将匹配结果可视化(见图3)。

图3 可视化结果
Part 3
计图性能对比

通过计图深度学习框架的实现与加速，superglue模型实现了显著的性能提升，同时消耗显存大幅下降，以下测试使用到的批大小为30, 对2000个描述子为128维的特征点进行匹配，使用的硬件设备为RTX 3090，下表给出Jittor，Pytorch，TensorRT点云匹配运行速度与显存消耗对比。
表1 点云匹配运行速度与显存消耗的对比
这里，Jittor使用了一个关键优化——Jittor自动混合精度计算，在以往PyTorch框架中，用户往往需要手工将类型转换成半精度浮点数，才能够实现性能加速，为了进一步提高易用性，计图同时提供了两种接口，为易用性和可定制性提供了多种选择：
# 手工类型转换model.float16()# 自动类型转换with jt.flag_scope(auto_mixed_precision_level=5):...
其中，我们为多种混合精度模式提供了不同级别的优化开关：
  • 第0级：默认使用float32
  • 第1级：默认使用float32,但开启tensorcore
  • 第2级：同第1级，但前向优化更加激进
  • 第3级：同第2级，前向反向优化都更加激进
  • 第4级：全局使用float16, 但是部分算子使用float32（如求和等）
  • 第5级：同第4级，并且参数也会使用float16
  • 第6级：全局所有算子强制使用float16

欢迎大家使用计图，同时也欢迎加入我们社区进行交流讨论，欢迎为我们提出建议，或者点亮star. 

Jittor github：https://github.com/Jittor/jittor

提出issue：https://github.com/Jittor/jittor/issues

Jittor QQ 群：761222083

计图论坛：https://discuss.jittor.org/

       

参考文献

  1. 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.

  2. Rocco I, Cimpoi M, Arandjelović R, et al, Neighbourhood consensus networks, Advances in neural information processing systems, 2018, 31.

GGC往期回顾

1. CVMJ刊发视觉注意力机制综述，Github 仓库获star1400+

2. Jittor团队提出大核注意力，开源新的骨干网络VAN

3. 第十届CVM会议4月在北京举行，会议议程发布了！

4. 首届JDC召开，计图技术委员会评出2021年度优秀Jittor开源项目

5. CVMJ首次入选，列中科院期刊分区表二区

您可通过下方二维码，关注清华大学计算机系图形学实验室，了解计算机图形学、Jittor框架、CVMJ期刊及会议的相关资讯。
 

Search

    Table of Contents