对标Theano、Caffe，清华开源首个国内高校自研深度学习框架——Jittor

对标Theano、Caffe，清华开源首个国内高校自研深度学习框架——Jittor

近日，清华大学计算机系图形实验室宣布开源一个全新的深度学习框架：Jittor，中文名计图，对标Theano、Caffe。

据悉，计图（Jittor：Just in Time）是一个采用元算子表达神经网络计算单元、完全基于动态编译（Just-in-Time）的深度学习框架，元算子和Numpy一样易于使用，并且超越Numpy能够实现更复杂更高效的操作。而统一计算图则是融合了静态计算图和动态计算图的诸多优点，在易于使用的同时，提供高性能的优化。基于元算子开发的深度学习模型，可以被计图实时的自动优化并且运行在指定的硬件上，如CPU，GPU。

Jittor特性

元算子融合

我们将神经网络所需的基本算子定义为元算子，多个元算子之间可以相互融合成更加复杂的算子，如卷积层，归一化层，进-步构成神经网络和深度学习应用。

高阶导数以及反向传播闭包

元算子是反向传播闭包，既元算子的反向传播也是元算子。同时支持计算任意高阶导数。在深度学习算子开发过程中，免去了反向传播算子重复开发工作，同时可以使用统- -的优化策略。

算子动态编译

Jittor内置元算子编译器，可以将用户通过元算子编写的Python代码，动态编译成高性能的C+ +代码。

自动优化

Jittor内置优化编译遍(complier pass)，同时和LVM兼容，这些编译遍会根据硬件设备，自动优化动态编译的代码，常见的优化编译遍有:循环重排，循环分裂，循环融合，数据打包，向量化，GPU并行。这些编译遍对C+ +代码进-步优化， 生成对计算设备友好的底层算子。

统一内存管理

Jittor使用了统- -内存管理，统- -GPU和CPU之间的内存。当深度学习模型将GPU内存资源耗尽时，将使用CPU内存来弥补。

高效同步异步接口

同步接口编程简单，异步接口更加高效，Jittor同时提供这两种接口，同步和异步接口之间的切换不会产生任何性能损失，让用户同时享受到易用性和高效率。

模型迁移

Jittor采用和PyTorch较为相似的模块化接口，为了方便用户上手，Jittor提供了辅助转换脚本，可以帮助您将PyTorch的代码转换成ittor的模型。在参 数保存和数据传输上，Jittor使用和PyTorch-样的Numpy+ pickle协议,所以Jittor和PyTorch的模型可以相互加载和调用。

下面的代码演示了如何一步一步使用 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()}")

最后附上官网地址：https://github.com/Jittor/jittor