清华大学计算机系教授胡事民：自研深度学习框架“计图”2 大创新、6 大特性详解

6 月 21 日，由北京智源人工智能研究院主办的 2020 北京智源大会正式开幕（直播入口： https://2020.baai.ac.cn ），大会为期四天，各主题论坛和分论坛将围绕如何构建多学科开放协同的创新体系、如何推进人工智能与经济社会发展深度融合、如何建立人工智能安全可控的治理体系、如何与各国携手开展重大共性挑战的研究与合作等一系列当下最受关注的问题进行交流和探讨。

在智源大会第三天（6 月 23 日)下午的 AI 框架主题论坛上，清华大学计算机系教授胡事民作了题为《统一计算图：机器学习框架「计图」的创新与探索》的演讲。他分享了国产深度学习框架“计图”所做的创新与探索，详细介绍了该框架的特性及创新点，及最新研究进展。

以下内容根据胡事民的演讲整理，未经本人确认。

为什么要研发国产深度学习框架“计图”？

演讲中，胡事民谈到了深度学习框架“计图”研发的初衷。

他表示，回溯深度学习框架在过去十年间的演进，国外的 TensorFlow 、PyTorch 、caffe 等主流深度学习框架发展迅速，国产深度学习框架的发展相对慢一些。

我国在人工智能领域取得了重要进展但发展不均衡，我们长于算法，弱于学习平台。我国的 AI 研究和算法研究多基于 TensorFlow 、PyTorch 等机器学习平台，这容易陷入“卡脖子”困境。

建设机器学习平台难度很大，需要对机器学习算法有深刻理解，精准把握图形图像应用，及巧妙应用底层系统软件。尽管挑战大，但中国应该积极自主研发机器学习平台，这既是基于自主创新的需要，也是基于人才培养的需要。

也正是在这样的背景下，基于元算子和统一计算图的深度学习框架“计图”（Jittor）应运而生。2020 年 3 月 20 日，计图正式发布。

胡事民表示，计图带来一次深度学习框架跨越的机遇，他希望该框架能够培养国内机器学习平台人才，解决对国外平台依赖性的问题，希望计图有机会能够超越国外的主流框架。

“计图”架构全景图及六大特性

据胡事民介绍，计图的整体架构自下而上分为四层：系统层、算子层、框架层、应用层，具体如下图：

在多种应用中，相比国际主流的平台，计图的性能提升了 5%-153% 。

胡事民表示，计图框架具有六大特性：

1、反向传播闭包

元算子是反向传播闭包，元算子融合的算子也是反向传播闭包，使 Jittor 可以统一“前向”和“反向”计算图；并可求高阶导数。

2、算子动态编译，运行时优化

用户编写的算子与模型，将在运行的时候，生成高性能的代码，并且进行动态编译优化。

3、统一内存管理，节省内存

如果 GPU 内存耗尽了可以使用 CPU 内存弥补

4、极简的自定义算子开发（提出 Code 算子）

用户可以在 Python 中内联 C++代码，数行代码即可完成高性能算子的开发。

5、辅助模型迁移工具

能做到“一键迁移”，迁移工具能帮助用户快速地从 PyTorch 迁移到 Jittor 。

6、计图模型库

今年 5 月 30 号，计图发布了 GAN 模型库。胡事民透露，将在 6 月底发布计图分割库。此外检测库，3D 点云库，3D 网格库等更多模型库将陆续发布。

两大创新点

胡事民表示，计图的创新点主要体现在两个地方：一是统一计算图。静态图高效，动态图易用性好，但二者却不可兼得，针对这个行业痛点，统一计算图是一个好的解决方案。

第二个创新点是元算子，针对业内存在的算子开发和优化困难痛点，元算子融合与动态编译优化能很好解决这个问题。

接下来，胡事民重点介绍了这两大创新点。

元算子融合

该创新特性设计主要针对行业痛点出发。Tensorflow 算子数量 2000+，PyTorch 算子数量 700+，庞大的算子库维护、优化困难。

计图创新性的将神经网络计算所需的基本算子进行了归纳总结，提出了元算子（18 个）和算子融合的概念。具有效率高，易于开发，可以统一优化。

元算子一共分为三类：重索引、重索引化简、元素级。

其中，重索引算子是一类一对多映射关系的元算子。常用的重索引算子有广播（Broadcast)、填补（Pad）、切分（Slice）。

重索引化简算子是一类多对一映射关系的元算子。常用的重索引化简算子有累乘（Product) 、累加（Sum)、取平均值（mean )。

元素级算子是一类一对一映射关系的元算子。逐个元素的运算都属于元素级元算子。

元算子是反向传播闭包（元算子的反向传播也是元算子)。云算子可以自动生成其反向传播算子。

多个元算子之间可以进行相互融合。元算子可以通过融合算法，自动生成复杂的算子。

元算子具有完备性，它可以覆盖绝大多数深度学习算子。对于元算子无法覆盖的算子，计图提供了 Code 算子，以此完成 100% 的算子覆盖。

统一计算图

计算图是所有深度学习框架用来描述模型的数据结构。计算图可分为静态计算图和动态计算图。

静态计算图和动态计算图各有优缺点，静态计算图的优点是高效，缺点是灵活性差；动态计算图的优点是易用，灵活性高，缺点是效率高。

易用性和高效性不可兼得。对此，计图做出了“统一计算图”的创新。

统一计算图兼顾高效和易用，除了统一静态图和动态图，统一计算图还完成了多种统一的计算图：

• 统一管理前向反向图，支持高阶导数；传统机器学习框架区分前向和反向传播，统一计算图将多次迭代的前向和反向计算图拼接，节省了大量的重复计算，提升灵活性和优化空间。

• 统一管理 CPU-GPU 内存，突破 GPU 显存限制统一；

• 统一同步异步运行接口，使得数据读取，内存拷贝，模型计算可以同时进行，提升性能。

• 统一管理多次迭代的计算图，使得框架可以实现跨迭代的融合优化。

计图的特点与新进展

1、动态编译

计图内置的元算子编译器，可将 Python 代码动态编译高性能 C++代码。

计图内置 LLVM 兼容的优化编译遍（Pass)。这些编译遍会根据硬件设备，自动优化动态编译的代码，并对 C++代码进一步优化，生成对计算设备友好的底层算子。

2、骨干网络

计图目前支持包括 AlexNet、VGG、ResNet 等在内的 30+骨干网络，并还在不断完善中。

3、GAN 模型库

计图目前支持包括 27 个主流的 GAN 模型库，包括图像分类、检测、分割、机器疼强化学习等。在性能方面，Gan 模型库速度比 PyTorch 均得到提升，半数模型性能提升超过 1 倍。

4、模型转换工具

计图和 PyTorch 的接口较为相似，内置自动转化脚本，可轻松实现转化。

5、分布式。

无需修改训练代码，只需修改启动命令；自动数据并行，实现多卡训练。

演讲中，胡事民带来了计图的最新进展。首先在 Res2Net 神经网络骨干模型上，通过最新构建层间多尺度，实现更强特征表达，显著提升多种常见视觉分析模型性能，计图上得到 10%的性能提升。

此外在 PointNet ++ 上，PointNet ++ 分类准确率超过了原论文（92.3 VS91.9），模型代码已经公开在 jittor-online-first 上，更多几何学习模型和算法也将在后续不断发布。