清华开源自研深度学习框架

AI框架，又来重磅中国玩家。

刚刚，清华自研的深度学习框架，正式对外开源。“贵系”计算机系的图形实验室出品，取名Jittor，中文名计图。

值得一提的是，这也是首个来自中国高校科研机构的开源深度学习框架，之前，业内来自“高校”的还有加拿大蒙特利尔大学的Theano，UC伯克利的Caffe。

与主流的深度学习框架TensorFlow、Pytorch不同，Jittor是一个完全基于动态编译（Just-in-time）、使用元算子和统一计算图的深度学习框架。

研发团队介绍称，开发Jittor是为了将新技术、硬件和模型的能力，更好地释放出来。

“深度学习发展迅猛，TensorFlow、PyTorch这些老牌主流框架，也会在新模型，新算法，新硬件上表现不佳，所以需要新的框架，在易于扩展同时保持高效。”

而现在框架呈现出来的能力，的确有超越“前辈”的倾向：

基于Jittor开发的深度学习模型，可以实时自动优化并运行在指定的硬件上，如CPU，GPU，在多种机器视觉任务上能够比同类产品PyTorch性能提高10%~50%。

团队还介绍，如此成果，得益于Jittor的两大创新点：元算子和统一计算图。这也是Jittor的立身之本。

Jittor的核心：元算子与统一计算图

根据Jittor官方文档定义，元算子是指神经网络所需的基本算子。

在TensorFlow，PyTorch框架底层，有上千个算子，如此多的算子使得开发和优化难度大幅提升。

在设计Jittor的时候，他们就定下了一个目标，即用户只需要数行代码，就可定义新的算子和模型。同时在保证易用的同时，不丧失任何可定制性。

所以在Jittor中，多个元算子之间，可以相互融合成更加复杂的算子，这些复杂算子构成了神经网络计算的多个模块，如卷积层，归一化层等等。

他们将这种融合称为元算子融合，可以提升性能，节省资源。在文档中，他们分享了一个案例：只用4个元算子，就实现了卷积操作。

他们介绍称，元算子的可拓展性很强，通过对元算子的简单修改，就可以实现更多复杂的卷积操作，如扩张卷积、深度卷积、点卷积、分离式卷积、反卷积等。

而且，通过元算子反向传播闭包，能自动生成反向卷积层。具体如下图所示，反向卷积层将来自输出的梯度，通过4个元算子，将梯度反向传播给卷积层的权重：

Jittor开发团队介绍称，在这样的设计下，元算子和Numpy一样易于使用，并且超越Numpy能够实现更复杂更高效的操作。

而且，通过元算子的反向传播闭包，Jittor可以对所有前向反向算子进行统一管理，这就是他们所说的第二个创新点：统一计算图。

简单来说，统一计算图是完成了多种统一的动态计算图。根据官方文档介绍，在Jittor中，核心有四个方面的统一：

• 统一管理前向反向计算图，使得高阶导数可以被支持。
• 统一管理CPU，GPU内存，使得训练模型时，可以突破原有的GPU显存限制，让CPU，GPU可以共享内存。
• 统一同步、异步运行接口，使得数据读取，内存拷贝，模型计算可以同时进行，提升性能
• 统一管理多次迭代的计算图，使得平台可以实现跨迭代的融合优化。

基于这个方面，他们给出了Jittor与其他各个框架的特性对比：

自动微分、动态图方面，Tensorflow、Pytorch和Jittor都支持。但在同步接口和异步接口方面，Jittor的优异性得到了体现。

同步接口易于编程，异步接口有助于提高性能，Jittor同时支持这两种接口。

相比之下，Tensorflow部分算子支持统一内存管理，而PyTorch不支持异步接口，而Jittor的所有算子都支持统一内存管理，当深度学习模型将GPU内存资源耗尽时，将使用CPU内存来弥补。

除此之外，Jittor还支持跨迭代融合。

在这些特性的支持下，Jittor具备了动态编译的能力。

官方文档介绍称，通过内置元算子编译器，可以将用户用元算子编写的Python代码，动态编译成高性能的C++代码。

比如，下图中的Python代码编写了神经网络中常用的批归一化层（batch norm）， 通过元算子编译器，动态生成了批归一化层C++代码。

开发团队介绍称，Jittor还会使用内置的编译优化，以及LLVM兼容的优化编译遍（complier pass）来优化动态生成的代码。

这些编译会根据硬件设备，自动优化动态编译的代码，常见的优化编译有：循环重排，循环分裂，循环融合，数据打包，向量化，GPU并行。

他们说，这些编译遍，能够对C++代码进一步优化，生成对计算设备友好的底层算子，从而提高性能。

这体现了他们设计Jittor的另一个理念：

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

此外，在整体设计中，他们还遵循了实现与优化分离的理念。

如此打造出来的整体架构，“用户可以通过前端接口专注于实现，而实现自动被后端优化。从而提升前端代码的可读性，以及后端优化的鲁棒性和可重用性。”他们介绍称。

Jittor的整体架构与上手样例

具体来说，Jittor的整体架构一共分为四层，如下图所示：

它是基于Jit编译技术、完全重新设计的深度学习框架，从上到下分别是应用层，前端层，后端层，硬件层，官方文档的介绍如下：

应用层的代码，用户使用Python编写，并可以访问从前端层公开的所有接口。

前端层 是Jittor的组件之一，代码用Python编写，提供了元算子的调用接口，来操作Jittor变量和Jittor实现的通用模型。

后端层是Jittor的内核，由C++编写，管理底层硬件资源。该层包含很多模块，比如算子融合器、第三方算子、JIT编译器、统一计算图、统一内存调度、统一执行器等。

硬件层支持的硬件有CPU和Nvidia GPU。但如果需要让Jittor支持新的硬件，只需要重载编译接口即可，让Jittor移植到新的硬件的难度将大大降低。Jittor开发团队说，他们将在未来支持更多的计算设备。

如此架构，用起来怎样？官方文档介绍称，从头只需要若干行代码，就能训练一个两层神经网络。

上面的代码，定义了激活函数和全连接层。Jittor开发团队介绍称，其实这些层已经集成在了框架中，并使用了类似的实现方式，在这里重新定义，用于更好展示内部机制和实现。

从代码中可以看出，Jittor的接口和现在主流深度学习框架接口类似，都是使用模块化的方式定义模型。其中，random、matmul、exp都是Jittor的内置算子。

基于JIT编译，Jittor的后端会将这几个算子自动融合成一个算子。

上面的代码，定义了双层神经网络。隐层的神经元个数是10， 使用的激活函数是上面定义好的sigmoid。

最后，可以从头开始训练模型。在这段代码，使用了梯度下降和L2 loss来训练网络。训练过程是异步的。

Jittor开发团队介绍称，Jittor会自动计算梯度并且将计算图保存起来，后端的JIT编译器会根据计算图，同时使用算子级别优化和图级别的优化。

他们进一步解释称，在这一示例中，Jittor使用了以下几种优化：

算子融合：激活函数和loss函数会被融合在一起。

并行化：算子会自动并行化以提升性能和计算密集度，在现代多核CPU和GPU上十分有效。

并发：没有依赖关系的操作会被并发执行，比如内存拷贝和计算可以并发并相互重叠。

元算子与统一计算图加持，整体框架优化下，Jittor在一些任务上展现出了性能提升，在多种机器视觉任务上尤为明显。

多个视觉任务上，性能超过现有主流框架

Jittor开发团队提供了实验数据。在ImageNet数据集上，使用Resnet50模型，GPU图像分类任务性能比PyTorch相比，提升32%；CPU图像分类任务提升11%。

在CelebA数据集上，使用LSGAN模型，使用GPU处理图像生成任务，Jittor比PyTorch性能提升达51%。