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++代码。

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

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

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

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

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

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

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

发布日期：2020年3月20日

来源：量子位