[模型库] Jittor开源点云模型库,多种主流模型训练比Pytorch提升一倍以上

2020年10月10日,第二届清华大学计图深度学习研讨会成功举办,研讨会上,计图还发布了检测与分割模型库、点云模型库和可微渲染功能。

点云数据的获取与处理是图形学和三维视觉中的重要问题,在测绘、自动驾驶等方面有着广泛的应用。由于其巨大的应用场景,近年来,点云处理引起了相关技术人员的广泛关注,涌现了一批非常出色的工作。图1是点云数据及其分割的示例。

图1 a. 点云数据的示例

图1 b. 点云分割的示例

计图框架本次发布的点云模型库包括几种最有影响力的模型:PointNet、PointNet++、PointCNN、DGCNN 和PointConv ,支持分类和分割。

我们在ModelNet40 数据集上面进行了分类实验,实验结果表明,我们在效果上已经达到了论文中的点数,甚至超过论文中的性能(见表1),这是因为Jittor在实现中添加了一些论文中没有使用的数据增强方式。

表1 ModelNet40 数据集上的分类结果

Model Input Overall accuracy Paper Report
PointNet 1024 xyz 89.6 89.2
PointNet ++ 4096 xyz normal 92.3 91.9
PointCNN 1024 xyz 92.6 92.5
DGCNN 1024 xyz 92.9 92.9
PointConv 1024 xyz normal 92.4 92.5

同时,在训练速度上,Jittor与Pytorch 相比有着明显的提升,其中 PointNet ++、PointCNN和 PointConv 在训练速度上超出 pytorch一倍以上,具体的时间对比见图2。

图2 ModelNet40 上的训练时间的对比

另外,我们也在ShapeNet 上面进行了分割的实验。在训练速度上,Jittor与Pytorch 相比,也得到了提升,其中PointNet++提升85%。图3是一些点云模型的分割效果。

图3 点云分割的效果

欢迎在计图官网或GitHub上下载Jittor的点云模型库:

https://cg.cs.tsinghua.edu.cn/jittor/

利用Jittor统一计算图的优势,可以很方便地调整计算图,针对您的研究,优化或设计新模型,开展高水平研究!

使用计图点云库有任何问题,可以在Github上提交issue,或加入如下计图-点云微信群,和计图团队实时交流和讨论。

为什么有些模型提升很多,比如PiontNet++,而有些模型(如PointNet)提升很少?

主要是跟 PointNet 和 PointNet ++ 的结构有关 :

  1. PointNet 模型结构简单,只用到了一些简单的算子 (比如线性层、池化层),这些常见算子 Pytorch 已经优化的非常好了,所以 Jittor 的提升不明显。
  2. PointNet ++ 模型结构比 PointNet 复杂,用到了一些自定义的算子(最远点采样、最近邻采样),这些算子 Pytorch 并没有进行优化,我们使用 Code 算子对这些用到的自定义的算子进行了优化,取得了大幅的提升。