对标 PyTorch，清华团队推出自研 AI 框架“计图”

vlambda
2020-06-12

对标 PyTorch，清华团队推出自研 AI 框架“计图” | AI 技术生态论

「AI技术生态论」人物访谈栏目是CSDN发起的百万人学AI倡议下的重要组成部分。通过对AI生态专家、创业者、行业KOL的访谈，反映其对于行业的思考、未来趋势的判断、技术的实践，以及成长的经历。

2020年，CSDN将对1000+人物进行系列访谈，勾勒出AI生态最具影响力人物图谱及AI产业全景图。本文为「AI技术生态论」系列访谈第3期。

对标 PyTorch，清华团队推出自研 AI 框架“计图” | AI 技术生态论

作者 | Just

头图 | CSDN 下载自视觉中国

出品 | AI科技大本营（ID:rgznai100）

在机器学习框架领域，Pytorch、TensorFlow已分别成为目前学界和业界使用最广泛的两大实力玩家，而紧随其后的Keras、Caffe/Caffe2、MXNet等框架也因为自身的独特性受到相应开发者的喜爱。

如今，AI开源框架之争再添新入局者。

如你所知，随着深度学习新技术的出现，任务复杂度不断提高，由于架构设计和不断扩充等原因，导致系统复杂，架构优化和移植变得困难，新模型的实际性能还有待提升。

一支清华大学团队决定研发更加灵活高效的深度学习框架。他们于近日宣布开源Jittor（计图），采用元算子融合和动态编译技术，深度优化内存，有效提升了系统的运行性能和通用性，确保实现和优化分离，大幅提升应用开发的灵活性、可拓展性和可移植性。

Github地址：

https://github.com/Jittor/Jittor

Jittor的研发团队是清华大学计算机系的图形学实验室，负责人现为胡事民教授，长期从事可视媒体智能处理的研究。2006年-2015年间，实验室得到两期国家973计划项目的资助，在可视媒体的认知计算、机器学习、几何计算、智能算法等方面开展研究。

2009年起，实验室开展系统软件研究，逐步布局AI平台的研发。期间，实验室奠定了在图形图像应用、机器学习算法和底层系统软件三大优势，2018年，在胡事民教授带领下，以梁盾、杨国烨、杨国炜和周文洋等一批博士生为主力的团队开始搭建Jittor平台。Jittor团队核心开发成员梁盾告诉AI科技大本营（ID:rgznai100），2019年底他们完成了Jittor的基本功能，随后经过内部测试，于3月20日正式对外发布并开源。

作为一个采用元算子表达神经网络计算单元、完全基于动态编译（Just-in-Time）的深度学习框架，Jittor在三大新的设计理念下进行开发：

1.易用且可定制：用户只需要数行代码，就可定义新的算子和模型，在易用的同时，不丧失任何可定制性。

深度学习采用的卷积神经网络是由算子（Operator）组成的一个计算网络，当前深度学习框架有多达2000种算子。Jittor将算子运算进一步分解，形成了更加底层的三类20余种元算子闭包，目前神经网络常用算子均可以使用元算子的组合进行表达。

什么是元算子？梁盾解释，元算子是Jittor的基本算子，主要包括三类：重索引算子（如填补、切分等），是一种一对多的关系；重索引化简算子（如累乘、累加等），是一种多对一的关系；元素级算子（如常见的元素级加减乘除等），是一种多对多的关系。

常见的神经网络算子，如卷积、池化、批归一化等操作可由元算子组合表达。元算子的提出主要是为了适应神经网络算子增长的需求，使得新算子的优化可以分解为元算子的融合和优化，避免单独对新算子优化，确保实现与优化的分离，提升可扩展性。

Jittor通过元算子融合实现深度神经网络模型

2.实现与优化分离：用户可以通过前端接口专注于实现，而实现自动被后端优化。从而提升前端代码的可读性，以及后端优化的鲁棒性和可重用性。

3.所有都是即时的：Jittor的所有代码都是即时编译并且运行，包括Jittor本身。用户可以随时对Jittor的所有代码进行修改，并且动态运行。

面向未来深度学习框架的发展趋势，Jittor利用元算子组合表达的优势，提出统一计算图进行优化，并从底层开始设计了一个全新的动态编译架构。

Jittor团队称，该架构支持多种编译器，确保实现和优化分离，大幅提升了应用开发灵活性、可拓展性和可移植性，与其他主流框架相比，具有多项先进特性。对标 PyTorch，清华团队推出自研 AI 框架“计图” | AI 技术生态论

这些领先的特性背后具体如何体现？

梁盾介绍，基于元算子组合表达神经网络的优势，Jittor提出统一计算图，融合静态计算图和动态计算图，提供高性能的计算：统一管理前向反向计算图，自动支持任意高阶导数的计算；统一调度CPU和GPU内存，支持超大模型的训练；统一同步异步运行接口，使得数据读取、内存拷贝、模型计算可以同时进行；统一管理多次迭代的计算图，实现跨迭代的融合优化。

基于此，Jittor团队在平台实现了ResNet、VGG、SSD、DeepLab、LSGAN等多个网络模型，根据他们提供的数据，与Pytorch相比，Jittor的推理和训练速度达到10%-50%的性能提升。

梁盾表示，Jittor在性能上的提升，主要得益于Jittor提出的元算子融合和统一计算图，优化计算，节省计算资源，提升访存效率。同时，Jittor设计的全新编译架构将元算子动态编译成高性能的C++/CUDA代码，并进一步通过与LLVM兼容的优化编译遍（complier pass），生成对计算设备友好的可执行代码。

无论是提到先进特性还是性能提升，从Jittor核心开发团队的官方口径，他们无疑要对标的是业内最主流的深度学习框架Pytorch。

Jittor开源消息发布后，开发者们在社交媒体平台上对其满是赞誉，不过要与Pytorch竞争，显然还有很长的路要走。

“中国作为人工智能产业发展和应用的最大市场，我们应该在人工智能生态的全产业链上占有一席之地。”

从国家科技战略层面，梁盾对AI科技大本营称，我国目前深度学习研究和应用主要依赖于国外平台，的确面临着卡脖子的风险，所以一些国内的IT企业也推出了自己的平台，希望Jittor和国内同行一起努力，为中国人工智能产业发展贡献一份力量。

最后，你怎么看Jittor？AI科技大本营将邀请业务专家做进一步解读，也欢迎在留言区发表你的观点。

vlambda博客
学习文章列表