清华大学开源首个国内高校自研深度学习框架Jittor

Posted 2021-04-19 深度学习初学者

深度学习技术正广泛应用于人工智能的各个领域，如计算机视觉、机器翻译、自然语言处理、智能机器人等，取得了前所未有的突破。当前，一方面，随着深度学习新技术的出现、任务复杂度的提高，易于扩展同时保持高效的架构成为发展趋势；另一方面，我国人工智能产业发展迅速，急需构建自己的开源深度学习生态。

清华大学计算机系胡事民教授研究团队提出了一个全新的深度学习框架——计图（Jittor）。Jittor是一个采用元算子表达神经网络计算单元、完全基于动态编译（Just-in-Time）的深度学习框架。

图1 “计图”通过元算子融合实现深度神经网络模型

深度学习采用的卷积神经网络是由算子（Operator）组成的一个计算网络。由于架构设计和不断扩充等原因，当前深度学习框架有多达2000种算子，系统复杂，优化和移植困难。Jittor则将算子运算进一步分解，形成了更加底层的三类20余种元算子闭包，目前神经网络常用算子均可以使用元算子的组合进行表达。面向未来深度学习框架的发展趋势，Jittor利用元算子组合表达的优势，提出统一计算图进行优化，并从底层开始设计了一个全新的动态编译架构。该架构支持多种编译器，实现了所有代码的即时编译和动态运行，确保了实现和优化分离，大幅提升了应用开发灵活性、可拓展性和可移植性。

图2 “计图”与其他平台的计算图特性对比

Jittor与国际主流平台相比，具有多项先进特性（图2）。目前ResNet、VGG、SSD、DeepLab、LSGAN等多个网络模型已经在Jittor平台实现，可供用户使用。与同类型框架相比，Jittor在收敛精度一致情况下，推理速度取得了10%-50%的性能提升（图3）。

图3 Jittor和PyTorch推理与训练速度对比

Jittor的研发得到了国家自然科学基金创新群体项目、北京信息科学与技术国家研究中心团队项目和清华-腾讯联合实验室项目的资助。期望Jittor能为学界和业界提供一个灵活高效的深度学习平台，促进人工智能的研究和应用，赋能人工智能产业。

Jittor 是一个基于即时编译和元算子的高性能深度学习框架，整个框架在即时编译的同时，还集成了强大的Op编译器和调优器，为您的模型生成定制化的高性能代码。

Jittor前端语言为Python。

前端使用了模块化的设计，类似于PyTorch，Keras，后端则使用高性能语言编写，如CUDA，C++。

下面的代码演示了如何一步一步使用Python代码，从头对一个双层神经网络建模。

import jittor as jtfrom jittor import Modulefrom jittor import nnclass Model(Module): def __init__(self): self.layer1 = nn.Linear(1, 10) self.relu = nn.Relu()  self.layer2 = nn.Linear(10, 1) def execute (self,x) : x = self.layer1(x) x = self.relu(x) x = self.layer2(x) return x
def get_data(n): # generate random data for training test. for i in range(n): x = np.random.rand(batch_size, 1) y = x*x yield jt.float32(x), jt.float32(y)
model = Model()learning_rate = 0.1optim = nn.SGD(model.parameters(), learning_rate)
for i,(x,y) in enumerate(get_data(n)): pred_y = model(x) loss = ((pred_y - y)**2) loss_mean = loss.mean() optim.step (loss_mean)    print(f"step {i}, loss = {loss_mean.data.sum()}")

目前 Jittor 已开源，用户可以采用 pip 等方法下载使用。 

Jittor 官网：https://cg.cs.tsinghua.edu.cn/jittor/  

项目地址：https://github.com/Jittor/jittor

了解更多信息，请访问“计图(Jittor)”官方网站：

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