在百度狂推PaddlePaddle框架，以及旷视计划本月25日发布自研深度学习训练框架MegEngine之际，清华大学突然发布首个中国高校自研深度学习框架Jittor，中文名为计图。
官网链接 https://cg.cs.tsinghua.edu.cn/jittor/
Github链接 https://github.com/Jittor/jittor
来源：AIZOO
如果把深度学习框架的发展比作春秋战国，那2018年前可以算作春秋时代，其中既有早期的Torch（2002年）、Theano（2007年），又有2012年深度学习的转折年之后的Caffe（2013年）、Deeplearning4J（2014年）、Keras（2015年）、TensorFlow（2015年）、MXNet（2015年）、日本创业公司的Chainer（2015年）、微软CNTK（2016年）、百度PaddlePaddle（2016年）、Caffe2（2017年）、PyTorch（2017年）等，甚至还有依附于这些框架的微框架——FastAI、TFLayer（Keras也算），种种深度学习框架层出不穷，烽烟四起。
但是，自2018年以后，深度学习可以说是步入了战国时期，强国开始不断侵蚀弱国的地盘，以至于今天形成了PyTorch和TensorFLow两强称霸的格局。曾经的霸主Caffe逐渐势微，但在工业界还有一席之位，Keras基本融入进了TensorFlow，Chainer宣布改旗易帜到PyTorch，Caffe2也融合进了PyTorch，MXNet虽然写的教材很棒，开源的模型也不少，但是仍旧不愠不火。另外，还有国产的PaddlePaddle，最近在各大公众号砸钱推广其Detection框架、Server部署框架、AI studio等，仍在向强国奋勇迈进。
旷视也预计本月25日发布其MegEngine训练框架，会同时开源多个模型示例，加入国产训练框架之争。没想到，半路突然杀出个程咬金，清华大学今日突然发布了国内首个高校研发的深度学习框架Jittor。
首先介绍一下Jittor的开发团队，该团队来自清华大学计算机系的图形学实验室，实验室负责人为胡事民教授，主要研发主力为实验室的梁盾、杨国烨、杨国炜和周文洋等博士生。
Jittor基本组成
根据Jittor官网的介绍，该框架是一个完全基于动态编译的框架，主要有两大核心部分：元算子和统一计算图。
元算子是基本的运算单元，也就是我们常说的OP，作者称元算子和Numpy一样易用，还超越了Numpy能够实现更复杂高效的操作。基本的元算子 可以融合成卷积、池化、BN等算子。下面是从元算子，到基本深度学习单元，再到模型的结构图。（有没有感觉这个图显得比较古老，有点类似10年前自动化系统的介绍图～）
Jittor基本结构组成
Jittor内置的元算子编译器，可以的将Python编写的算子，动态编译成C++代码。
统一计算图则融合了静态图和动态图的优点，易于使用，可以高效的将模型优化到CPU、GPU等硬件上。
Jittor内置的编译器，与LLVM兼容，会根据硬件，自动动态编译代码，对C++代码进一步优化，生成对底层硬件更友好的底层算子。
Jittor一个比较创新的地方是，其统一管理GPU显存和内存，当显存耗尽时，可以转到内存。这样大家就不用担心显卡的那区区8～12G显存不够用了。但是主板内存的效率，相比显存下降多少，需要测试比较。
Jittor支持显存和内存的统一管理
根据介绍，我们看到，相比各种国产的套个壳子的芯片或者操作系统，Jittor在底层还是有一定的创新的。
2. 前端使用
Jittor后端使用CUDA、C++编译，但前端是Python。前端则非常类似于PyTorch。由于参数保存和数据传输与PyTorch一样，都是用Numpy和Pickle格式，Jittor甚至可以直接加载PyTorch的模型。厉害了吧～
下面是一个两层的全连接网络的搭建和训练示例，相信会用PyTorch的同学，直接import jittor as torch就可以无缝切换使用了。这个真不是玩笑～
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，唯一的不同就是不需要loss.backward()，可以说真的基本无缝迁移啊。
Jittor还贴心的给出了PyTorch的转换脚本。笔者想说，这个转换真的是容易呀，虽然他们后台不一样，但是脸，简直就是双胞胎一对啊。所以大家从PyTorch转过来，简直没有学习成本。
PyTorch 2 Jittor
关于安装等信息，大家可以去官网查看，这里不再赘述。
3. 模型示例
可以说，一个框架能不能火，除框架本身的易用性外，其官方例子，以及开发者是否积极贡献代码，都对一个框架的发展有重要影响。这不，百度Paddle团队就疯狂的实现了各种模型，基本上只要TensorFlow有的，Paddle都会搞一套。
Jittor在上午也同步开源了SSD、DeepLabV3+、LSGAN的完整示例。
- 基本概念：Op, Var
- 元算子：通过元算子实现自己的卷积层
- 自定义算子：使用C++和CUDA编写您的算子，并其进行即时编译
- 示例1：线性回归
- 示例2：MNIST图像分类
- 示例3：图像生成之 LSGAN
- 示例4：语义分割之 DeepLabV3+
- 示例5：目标检测之 SSD
在其实现的VGG16作为Backbone的SSD中，在Pascal数据集上，Jittor的mAP甚至还高千分之一，其推理速度相比PyTorch要快10~50%.
Jittor推理速度比较
笔者大致看了一下其实现的DeepLabV3+和SSD，写的还是非常好的，初学者可以参考一下它的代码结构。
4 总结
总体来说，Jittor在后端还是下了一番功夫的，其前端也保持和PyTorch基本一模一样的接口，保证了前端的易用性，方便开发者迁移。
不过，在深度学习框架的战国时期，各个训练框架的基本功能单元都很相似，甚至有ONNX这种通用格式，或许训练框架会越来越走向大一统，现在看来PyTorch是风头正盛， TensorFlow也不甘人后，Paddle也在花重金推广。在部署方面，TensorRT、OpenVINO和TVM也基本形成三国鼎力之势，Jittor能否杀出重围，大家可以在评论区发表一下您的看法。
不管怎么说，国货当自强。虽然Caffe、MXNet、TVM都是中国人主导开发的，可惜都是在美国高校和公司开发的。我们祝愿Jittor、PaddlePaddle和一周后发布的旷视MegEngine能做大做强，甚至走出国门。
推荐阅读
关注元峰微信号“AIZOOTech”
更多人工智能资讯，算法模型相关请关注AIZOO专栏
