清华开源Jittor：首个国内高校自研深度学习框架，一键转换PyTorch
chenfenggang 2020-03-20
继 Theano、Caffe 之后，又一个由高校主导的深度学习框架开源了，而且还是国产。
深度学习框架越来越多，主导的团队也从高校研究机构渐渐转向了科技巨头。但是，学界在这一领域的力量不容忽视。今日。清华大学开发了一个名为计图（Jittor）的深度学习框架。这一框架有望为深度学习社区提供新的方案，也能够推动深度学习框架国产化的进程。
据悉，计图（Jittor：Just in Time）是一个采用元算子表达神经网络计算单元、完全基于动态编译（Just-in-Time）的深度学习框架，其主要特性为元算子和统一计算图。
据官网介绍说：「研究团队将神经网络所需的基本算子定义为元算子」，元算子非常底层，通过相互融合可以完成复杂的深度学习计算，体现出了易于使用的特点。并且研发团队表示，采用元算子的 Jittor 已超越 Numpy，能够实现更复杂更高效的操作。
而另一层面，统一计算图则是融合了静态计算图和动态计算图的诸多优点，在易于使用的同时，提供高性能的优化。基于元算子开发的深度学习模型，可以被计图实时地自动优化并且运行在指定的硬件上，如 CPU、GPU。
目前 Jittor 已开源，用户可以采用 pip 等方法下载使用。
Jittor 官网：https://cg.cs.tsinghua.edu.cn/jittor/
项目地址：https://github.com/Jittor/jittor
即时、易用、可定制：实现和优化分离的框架
据官网介绍，Jittor 在设计时秉持易用、灵活、即时的设计理念：
- 易用且可定制：只需要数行代码，就可定义新的算子和模型。
- 实现与优化分离：可以通过前端接口专注于实现，而实现自动被后端优化。
- 所有都是即时的：Jittor 的所有代码都是即时编译并且运行的，包括 Jittor 本身。用户可以随时对 Jittor 的所有代码进行修改，并且动态运行。
Jittor 作为全新的深度学习框架，有几大新特性值得关注。首先，Jittor 采用了元算子的概念，将各种基本计算定义为元算子，并通过结合不同元算子，实现深度学习中的各项算子功能。
元算子融合。
其次，在算子的设置上，团队将元算子的反向传播进行了闭包，即元算子的反向传播也是元算子。这样避免了重复开发。此外，还支持计算任意高阶导数。
高阶导数及反向传播闭包。
然后，在编程语言上，Jittor 采用了灵活而易用的 Python。用户可以使用它，编写元算子计算的 Python 代码，然后 Jittor
将其动态编译为 C++，实现高性能。
算子动态编译。
更重要的是，Jittor 内置了优化编译遍（complier pass）功能，和 LLVM 兼容。优化过程会根据运行的硬件而定，以下为已支持的优化编译遍。
自动优化。
在内存管理上，CPU 和 GPU 内存统一管理。GPU 内存不够的时候会调用 GPU。
统一内存管理。
而在接口方面，Jittor 同时提供了同步和异步接口，切换不会造成性能损失。从而带来了易用性和高效率。
高效同步异步接口。
最后值得注意的是，为了帮助 Jittor 的发展，并吸收优秀框架的特性。Jittor 采用了和 PyTorch 相似的模块化接口，并有辅助转换脚本，还有和 PyTorch 一样的 Numpy+pickle 协议，使得两者的模型可以互相加载和使用。
模型迁移。
有了如此众多的特性，估计用来写模型会有非常好的体验。此外，Jittor 使用 Python 和 C ++编写，也支持 GPU 与 CUDA。
测评结果
Jittor 的性能如何，官方也提供了测评结果进行参考。目前 ResNet、VGG、SSD、DeepLab、LSGAN 等多个网络模型已经在 Jittor 平台实现，可供用户使用。与同类型框架相比，Jittor 在收敛精度一致情况下，推理速度取得了 10%-50% 的性能提升。
示例代码
Jittor 前端语言为 Python。前端使用模块化的设计，类似于 PyTorch，Keras，后端则使用高性能语言编写，如 CUDA，C++。Jittor 官网目前已提供了示例代码，我们可以来分析一下。
下面的代码演示了如何一步一步使用 Python 代码，从头对一个双层神经网络建模。
import jittor as jt
from jittor import Module
from jittor import nn
class 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.1
optim = 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()}") 以上为实现一个简单的全连接神经网络的代码。可以看到，类似于 PyTorch，只需要导入包，采用类继承的方式定义模型、数据处理（分批等）方式，并设置训练循环即可。从上手难度来看，熟悉 PyTorch 的用户可以直接上手，甚至没有因为变量名不同而带来影响。
研究团队
据官网介绍，Jittor 的开发团队来自清华大学计算机系的图形学实验室，实验室负责人为胡事民教授。
