如何评价清华大学发布的自研深度学习框架-计图(Jittor)？

笨老孩

2020年3月20日，清华自研的深度学习框架，正式对外开源。清华大学计算机系的图形尝试室出品，取名Jittor，中文名计图。

日天日地

领域相关，也来凑个热闹。
看到这个工作还是比较开心的，关于AI技术，国内以及华人圈子已经在发起越来越多具备硬核属性的努力了。
硬件上，有海思的Da Vinci，燧原的IPU，中科院的寒武纪(以及我司的含光)等等，软件上，有TVM，头条&港大合作的BytePS，商汤的性能优化团队（除了不开源都挺好）的底层优化工作(说手工优化似乎不完全合适，还是说手工优化+半手工优化+AI编译优化比较合适^-^)，华为在AI底层软件方面的工作(在AI编译器方面，华为是国内最早投入的公司之一)，以及这次的Jittor的工作(也不要落了进辉兄的Oneflow工作，能够在AI Infra这种硬核方向坚持创业多年，是很值得尊敬的努力)，当然内举不避亲，我们阿里PAI团队在AI硬核方向也投入不少，分布式、资源优化、编译优化、模型压缩、强化学习框架、迁移学习平台、建模辅助工具等等，欢迎感兴趣的同学勾搭，我可以帮忙route到合适的团队里去（此处满满的招聘带货嫌疑^-^）。
回归正题。
作为一个探索性的工作，我觉得Jittor是值得尊敬的，能够在这个工作里看到不少好的创新点。因为我对底层关注相对多一些，所以直接举一些偏底层的features，包括:

• 统一内存
  
• 跨迭代融合
  
• 以及一套新的AI编译框架的尝试
  

统一内存我们去年在生产集群里出于集群资源优化的需要，也做了类似的feature(@孙敏敏  )，也在训练作业里，发现有一些用户比较喜欢这个feature(麻麻我的作业不用再担心显存不够用被钉，只是可能会变慢^-^)。
跨迭代融合在目前我所知道的AI框架里还不支持（当然可以通过对现有框架hack来达到这个效果），而跨迭代融合功能的缺失会导致一些inter-minibatch的优化机会的丢失，这一点我没记错的话，去年在BytePS的工作里似乎还专门提到过(P3里是不是有提到映象不深，但本质上是相关的)。
新AI编译框架的尝试，倒不用多说，看到code和sample就足够self-explaining了。能够自己定义一套DSL，基于这套DSL来描述算子生成的空间，并打通完整的通路，也具备了一定的fusion能力的支持，即便只是在一些示例模型上跑通，这个工作也还是值得尊敬和认可的。
总的来说，值得赞赏。接下来我也会表达一些其他的comments：

• 从技术创新性上来说，可能还有比较长的路要走。我目前对Jittor的观察，还是做了大量的系统engineering的工作的，如果能够把相较于一些existing框架的benchmark的性能差异breakdown到具体的系统创新点，可能会更有助于其他同学了解Jittor提供的差异化，哪些是系统实现层面的，哪些是设计原则层面的，这也可能更好扩大其impact；
  
• 算子Fusion有了一定的支持，目前还有些局限，应该还能够做得更好。比如目前的代码实现里，能够看到遇到reduce/broadcast之类不符合embarrassingly-parallel特性的指令，fusion就会中断，对于比较复杂的计算图，这会显著影响到fusion的颗粒度。如果Jittor的研发同学有兴趣，欢迎来一起交流，这方面我们做过一些工作 (@龙国平 )，在生产环境也取得了一定的效果。
  
• 关于schedule探索生成和算子描述分离，我对于Jittor团队自己重新设计一套具备一定普适性的computation描述语言的这个taste还是很欣赏的，我们在做AI编译优化工作的时候，曾经想做类似的事情，还是限于工业界面向生产环境的的压力，没有做得那么激进，还是选择了针对访存密集算子设计了一套简化的schedule描述规则，对于计算密集算子直接follow了TVM/Halide的体系。但同时我想point out的是，自己从头造轮子的代价是大量的工程建设工作，也就需要有比较陡峭的成长curve。比如以深度学习里常见的reduction操作为例，CUDA上至少有三种不同性能表现(atomic/warp shuffle/shared memory)，适合不同context的实现，在Jittor里怎样能够自然的描述出这个探索空间？以及，如果想支持TensorCore这样的特殊硬件计算单元，在计算描述上怎样cover？
  
• 关于JIT编译的内容。如果我理解没有错，在Jittor的设计理念里，考虑到了dynamic shape的问题（如果不准确请指出），这个我倒觉得蛮好的。因为JIT的一个挑战就是编译期overhead，因为和AoT不同，JIT的编译开销会直接体现在运行过程中，如果对dynamic shape的问题处理不当，很容易因为JIT编译开销过大吃没了其带来的性能收益。在这方面我们也结合生产环境的需要做了一些有意思的尝试，算是比较好的解决了这个问题，感兴趣也欢迎来交流：）
  
• 作为一个AI框架，需要考虑的东西很多，除了一些有意思的想法实现，跑通sample model流程的示例以外，还存在着大量的细节。分布式怎么支持？分布式里还有不同变种，Data-parallelism, Model-parallelism., Pipeline-parallelism分别怎么支持？训练和推理环节的有效打通(冗余节点的去除，计算图动态性语义的导出保证等等)怎样支持？如果在这个框架上进行一些模型优化工作（比如量化和剪枝等）怎样支持？至于说到相关的生态建设，更是海量工作了。
  
• 这里面一些好的想法，如果能够在一些existing工作上进行拓展，也许能够有更大的impact，我指的是既包括实际的应用impact，也包括research impact。比如关于统一内存的功能，加在TF或PyTorch的内核里，会是一个很nice的feature。跨迭代融合稍微tricky一些，对于强势社区，被accept会不太容易，不过我个人确实对于跨迭代融合能够带来的端到端收益会有一些保留观点（同样欢迎指正）。而JIT编译框架的尝试，其实我个人觉得Jitter的思想和TVM是有相通的地方，并不存在框架上根本性的矛盾，如果是我来操作，可能会更愿意在TVM里进行扩展(比如加一些contrib的computation/schedule的描述体系，我个人感觉 @陈天奇 是welcome这样的贡献的^-^)，这样能够和领域同行形成有效的碰撞，可能产出更有impact的成果。并且也许在做这个贡献的过程中，会发现有些问题已经被已有的工作解决得比较好了，那么就可以把精力花在一些更有意思，更有创新性的问题上了。
  

最后，还是很appreciate Jittor团队的工作，期待其未来能够走得更好。
也欢迎来勾搭/交流/合作，多多响应扬清回复里的号召：）。

小怡妞2

利益相关：Jittor开发者梁盾七年多的室友。
我在清华计算机系是个菜鸡，有大佬大神在轮不到我来评价Jittor的优缺点，这个答案只是想安利各位同行尝试一下这个框架，多给开发团队提意见。盾神是我认识的人里最有极客范儿的之一，热爱钻研自己感兴趣的技术并且享受实现的过程，不追热点不水论文不出风头，关注CS领域的最新技术然后去学习去看文档去看源码去实现，最后通过自己的努力做出真正实质性的贡献去改变，这整个过程就是他最大的乐趣来源。我一直认为这才是真正优秀的计算机系学生，既佩服又羡慕他。盾神高中就玩斯坦福兔子，大学拿超算世界冠军，读博又深入底层做了这样优秀的深度学习框架，但为人还是一如既往地低调。我以为盾神会用自己的知乎账号答这个题，没想到他们注册了一个叫Jittor的账号。对于这样的人，最开心的事莫过于自己辛辛苦苦做出来的东西能够被更多的人使用、能够有自己的使用者社区。而且盾神是我认识的同级清华计算机系博士生里最强的，我相信他花这么多时间精力出来的成果也不会让人失望，所以真的强烈推荐大家去试着用一下，毕竟从pytorch上switch过去也很方便。
盾神开发Jittor的这一年多是我见过的他工作最辛苦的一段时间，作为一个在清华计算机系被称为神的家伙以及前OI国家队候补队员，盾神和其他的神一样，对于我们系专业课的作业都是随便花点时间就拿高分的，而且盾神代码和工程能力很强，基本上看不到他有需要每天码很久连续码很多天的情况。在写Jittor之前，盾神在CV和机器人方面做过一些工作，对tf和pytorch钻研颇多，而且跟我们这些调包怪不同，他对底层乃至硬件上的机制也都足够了解。但是Jittor不是一个大作业也不是一篇随便水水的论文，在AI会议动辄几千上万篇投稿的今天，想要做出真正有用的东西需要付出巨大的努力，一个重新设计的深度学习框架需要学很多东西写很多代码，对于一个学生团队来说非常不容易。这一年多看盾神从C++写到binding写到汇编写到算子库Cache硬件各种优化再写到接口，在这个过程中甚至还顺便搞了个公众号写优化方面的blog。Jittor经过反复迭代测试今天终于发布，作为好朋友很为他高兴。


盾神迭代开发测试过程中还设计了问卷，调查开发者们对深度学习框架的需求和使用习惯，之所以Jittor在接口上跟pytorch如此相似也是因为大部分开发者更习惯这样。没错，我说的唯一一句话是填这个问卷超出了我的能力范围，可见我有多菜。但即使是这么菜的我，盾神也找我聊了很多次下一步框架应该怎么扩展，大家可能还有哪些需求，他是真的希望花了这么多时间精力做出来的东西能够实实在在地有用，所以我相信他们团队一定会重视每一条有意义的issue，后续也会推出更多好用的模型和接口，扩展Jittor的适用领域，再次推荐大家多尝试多关注。

飞翔天空的确

简单看了一下代码，非常有意思的一个项目。因为看得不深入，所以如果有观察错误的话，我想对作者先致以歉意。
总的来说：项目更加关注在如何进行计算图优化以及just in time compilation上面（所以叫jittor），并不是关注完整的端到端的框架设计（比如说建模前端等等），我觉得定位是比较清楚的，自动代码生成是现在大家都很关注的方向，在校的同学能够着手把这一套都做一次，值得点赞。
一些能看到的工程点：

• 实现了一个比较经典的DAG graph，以及在图上来做fusion和各种pass。
  
• 从op的实现上，选择了细粒度的op，例如bcast，reduce，等等，然后通过这种方式来形成meta op，比如说convolution：https://github.com/Jittor/jittor/blob/master/notebook/meta_op.src.md
  
  
  
  • 值得关注的一点是，在XLA的早期，也有过对于op粒度的探索，目前大家的一些结论是，常见的op，比如说convolution，gemm，如果用细粒度op来实现，然后这些细粒度op是在一个op graph当中来做jit的，对性能会是一个很大的挑战（除了在代码里面embed constant value，loop reordering等等）之外，很多关于计算的细节信息都丢失了，会对后面的fusion pass有很大的挑战。
    
  • 现在一般的自动编译框架选择的方式其实是选择两层IR，一层做计算图DAG，一层做数学表达（比如说bcast，reduce，最典型的是Halide）。可能值得看一看。
    
  
  
• 编译是通过route到binary call，比如说nvcc和gcc，然后读取编译器的输出来做的（https://github.com/Jittor/jittor/blob/85d9ccc004b83ad2d83769ffd1803fb038a17264/src/jit_compiler.cc#L20）。这个让我想起当年的theano的设计，背后的思想一脉相承。相比较于对接libllvm等runtime library，是一个比较短平快可以来实现从前端到编译的办法。
  
  
  
  • 一个小tip，c++其实有demangle的一些utility（https://github.com/pytorch/pytorch/blob/master/c10/util/Type.h#L13），可以不需要把symbol的magic string放在代码里面（https://github.com/Jittor/jittor/blob/85d9ccc004b83ad2d83769ffd1803fb038a17264/src/jit_compiler.cc#L59）
    
  
  
• 因为编译的产出是直接的binary，不需要一个通用的framework runtime，所以编译以后的产出binary size会很小，这个也许会在端上设备、microcontroller等场景比较有意思？
  

很赞的一些工程细节：

• 有log和flag的设计（应该是学习了google glog和gflags？），并且代码当中做了比较多的logging，这个是个工程实现当中很重要的环节。https://github.com/Jittor/jittor/tree/master/src/utils
  
• 因为compilation时间比较久，所以考虑了做cache：https://github.com/Jittor/jittor/blob/master/src/utils/cache_compile.h
  
• 一个小亮点：通过term的special character来设置颜色的代码 https://github.com/Jittor/jittor/blob/85d9ccc004b83ad2d83769ffd1803fb038a17264/src/utils/log.h#L45 让人想起BBS灌水写签名档的年代，哈哈。
  

从水木盗的图

最后再次表示：非常有意思的一个项目，可以看到作者在学校关注最近的科研方向，同时又能很hands on的自己造一把轮子（褒义）的能力。
打个广告（团队同学已经比我手快了）：欢迎对机器学习、大数据、分布式系统等技术有兴趣的同学来一起实习、工作，或者在开源社区合作，夏季实习请扫码 :)
http://alibaba.tupu360.com/campusActivity/getActivityInfo?activityCode=nrIKmwibjn1bq9if1rSA6RiT3q5JjhsEdljYcLGSmuY%3D&enter=menu (二维码自动识别)

bribousensori

感谢各位专家的评价～不胜荣幸
刚刚发布收到了很多小伙伴的关注和疑问，我们研发团队也补充一下。
我先来说一下背景，现在深度学习框架主流的框架包括tf，pytorch，国内现在是深度学习研究的主力，但是国内急需更好的深度学习的生态，希望我们开源的深度学习框架能像星星之火可以燎原，也希望能够和学界和产业界共建和谐生态。
从我们自己官方网站 Jittor: 即时编译深度学习框架 中把设计理念抄过来：
易用且可定制 用户只需要数行代码，就可定义新的算子和模型，在易用的同时，不丧失任何可定制性。
实现与优化分离 用户可以通过前端接口专注于实现，而实现自动被后端优化。从而提升前端代码的可读性，以及后端优化的鲁棒性和可重用性。
所有都是即时的 Jittor的所有代码都是即时编译并且运行的，包括Jittor本身。用户可以随时对Jittor的所有代码进行修改，并且动态运行。
期望Jittor能为学界和业界提供一个灵活高效的深度学习平台，促进人工智能的研究和应用，赋能人工智能产业。
同时，Jittor才刚刚开始，前面的路还很长，平台的开发需要对图形图像应用的精准把握、对机器学习算法的深刻理解、对底层系统软件的巧妙应用。Jittor可能还存在一些错误和问题，并需要不断完善和演进，我们会一直努力，并欢迎批评和各种建议。
<hr/>感谢大家对Jittor的关注和评论，对于很多小伙伴提到关于学习成本的问题，Jittor在设计时其实也做了考虑，并进行了用户调研，最终决定使用与pytorch类似的接口，Jittor在方法名称和参数的设计时也尽量和pytorch保持一致。我们已经提供了API文档，对于简单的应用，我们相信有机器学习经验的小伙伴，能够轻松上手jittor！
<hr/>欢迎机器学习的爱好者加入我们。
了解更多信息，请访问“计图（Jittor）”
欢迎大家去github提交issue还有讨论：Jittor/jittor
官方网站。https://cg.cs.tsinghua.edu.cn/jittor。
Email：jittor@qq.com。
QQ群：761222083
<hr/>最新动态：
Jittor：新版本Jittor1.2.2推出JRender可微渲染库、支持Vision Transformer

木槿hr1

挺多朋友给我发这个项目，我粗略研究了一下，简单单理解是：
1，这个项目基本上是和xla, tvm, plaidml这样的编译器是更对等的，和pytorch, tensorflow 这样的框架反而距离远点，当然其实深度学习编译器和框架其实都在融合过程中。
2，在前端提供element, broadcast,reduce等基础操作来拼装任何其它高层次操作（和xla, tvm，plaidml都不同），内部全部转化成for loop优化，和tvm一样没去像polyhedral 处理一个更复杂的问题。前端基本上分几个类型吧，tvm, plaidml 那种类似DSL的描述方式；或者像XLA就是用一些固定集合的基本op来拼接。jittor 在做法上像xla，定义了一些基本op，但这些基本op（broadcast, reduce, element-wise等） 更底层，可以实现像tvm, plaidml那样的灵活性。当然要使用这些基本的op时，需要想清楚reindex怎么写，这不是很直接，需要费点脑细胞。
3，内部转化成for-loop之后，就可以做各种各样的优化操作了（各种pass)。
项目的技术含量还是蛮高的，思路很新颖，很高兴看到国内高校研究组能做出这样的成果，清华厉害。
2020/03/21 补充一些关于前端的思考
袁进辉：关于深度学习编译器前端的思考