大佬评清华Jittor,这是怎样一个深度学习框架?

590 阅读15分钟

点击上方“机器学习与生成对抗网络”,关注"星标"

获取有趣、好玩的前沿干货!

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

2020年3月20日,清华自研的深度学习框架,正式对外开源。清华大学计算机系的图形实验室出品,取名Jittor,中文名计图。——知乎问题

www.zhihu.com/question/38…


作者:贾扬清 

链接:www.zhihu.com/question/38…

简单看了一下代码,非常有意思的一个项目。因为看得不深入,所以如果有观察错误的话,我想对作者先致以歉意。

总的来说:项目更加关注在如何进行计算图优化以及just in time compilation上面(所以叫jittor),并不是关注完整的端到端的框架设计(比如说建模前端等等),我觉得定位是比较清楚的,自动代码生成是现在大家都很关注的方向,在校的同学能够着手把这一套都做一次,值得点赞。

一些能看到的工程点:

  • 实现了一个比较经典的DAG graph,以及在图上来做fusion和各种pass。

  • 从op的实现上,选择了细粒度的op,例如bcast,reduce,等等,然后通过这种方式来形成meta op,比如说convolution:github.com/Jittor/jitt…

  • 值得关注的一点是,在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,然后读取编译器的输出来做的(github.com/Jittor/jitt… library,是一个比较短平快可以来实现从前端到编译的办法。

  • 一个小tip,c++其实有demangle的一些utility(github.com/pytorch/pyt… string放在代码里面(github.com/Jittor/jitt…

  • 因为编译的产出是直接的binary,不需要一个通用的framework runtime,所以编译以后的产出binary size会很小,这个也许会在端上设备、microcontroller等场景比较有意思?

很赞的一些工程细节:

  • 有log和flag的设计(应该是学习了google glog和gflags?),并且代码当中做了比较多的logging,这个是个工程实现当中很重要的环节。github.com/Jittor/jitt…
  • 因为compilation时间比较久,所以考虑了做cache:github.com/Jittor/jitt…
  • 一个小亮点:通过term的special character来设置颜色的代码 github.com/Jittor/jitt… 让人想起BBS灌水写签名档的年代,哈哈。

从水木盗的图

最后再次表示:非常有意思的一个项目,可以看到作者在学校关注最近的科研方向,同时又能很hands on的自己造一把轮子(褒义)的能力。


作者:杨军 

链接:www.zhihu.com/question/38…

领域相关,也来凑个热闹。看到这个工作还是比较开心的,关于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:

  1. 从技术创新性上来说,可能还有比较长的路要走。我目前对Jittor的观察,还是做了大量的系统engineering的工作的,如果能够把相较于一些existing框架的benchmark的性能差异breakdown到具体的系统创新点,可能会更有助于其他同学了解Jittor提供的差异化,哪些是系统实现层面的,哪些是设计原则层面的,这也可能更好扩大其impact;

  2. 算子Fusion有了一定的支持,目前还有些局限,应该还能够做得更好。比如目前的代码实现里,能够看到遇到reduce/broadcast之类不符合embarrassingly-parallel特性的指令,fusion就会中断,对于比较复杂的计算图,这会显著影响到fusion的颗粒度。如果Jittor的研发同学有兴趣,欢迎来一起交流,这方面我们做过一些工作 (@龙国平 ),在生产环境也取得了一定的效果。

  3. 关于schedule探索生成和算子描述分离,我对于Jittor团队自己重新设计一套具备一定普适性的computation描述语言的这个taste还是很欣赏的,我们在做AI编译优化工作的时候,曾经想做类似的事情,还是限于工业界面向生产环境的的压力,没有做得那么激进,还是选择了针对访存密集算子设计了一套简化的schedule描述规则,对于计算密集算子直接follow了TVM/Halide的体系。但同时我想point out的是,自己从头造轮子的代价是大量的工程建设工作,也就需要有比较陡峭的成长curve。比如以深度学习里常见的reduction操作为例,CUDA上至少有三种不同性能表现(atomic/warp shuffle/shared memory),适合不同context的实现,在Jittor里怎样能够自然的描述出这个探索空间?以及,如果想支持TensorCore这样的特殊硬件计算单元,在计算描述上怎样cover?

  4. 关于JIT编译的内容。如果我理解没有错,在Jittor的设计理念里,考虑到了dynamic shape的问题(如果不准确请指出),这个我倒觉得蛮好的。因为JIT的一个挑战就是编译期overhead,因为和AoT不同,JIT的编译开销会直接体现在运行过程中,如果对dynamic shape的问题处理不当,很容易因为JIT编译开销过大吃没了其带来的性能收益。在这方面我们也结合生产环境的需要做了一些有意思的尝试,算是比较好的解决了这个问题,感兴趣也欢迎来交流:)

  5. 作为一个AI框架,需要考虑的东西很多,除了一些有意思的想法实现,跑通sample model流程的示例以外,还存在着大量的细节。分布式怎么支持?分布式里还有不同变种,Data-parallelism, Model-parallelism., Pipeline-parallelism分别怎么支持?训练和推理环节的有效打通(冗余节点的去除,计算图动态性语义的导出保证等等)怎样支持?如果在这个框架上进行一些模型优化工作(比如量化和剪枝等)怎样支持?至于说到相关的生态建设,更是海量工作了。

  6. 这里面一些好的想法,如果能够在一些existing工作上进行拓展,也许能够有更大的impact,我指的是既包括实际的应用impact,也包括research impact。比如关于统一内存的功能,加在TF或PyTorch的内核里,会是一个很nice的feature。跨迭代融合稍微tricky一些,对于强势社区,被accept会不太容易,不过我个人确实对于跨迭代融合能够带来的端到端收益会有一些保留观点(同样欢迎指正)。而JIT编译框架的尝试,其实我个人觉得Jitter的思想和TVM是有相通的地方,并不存在框架上根本性的矛盾,如果是我来操作,可能会更愿意在TVM里进行扩展(比如加一些contrib的computation/schedule的描述体系,我个人感觉

    @陈天奇

    是welcome这样的贡献的^-^),这样能够和领域同行形成有效的碰撞,可能产出更有impact的成果。并且也许在做这个贡献的过程中,会发现有些问题已经被已有的工作解决得比较好了,那么就可以把精力花在一些更有意思,更有创新性的问题上了。

最后,还是很appreciate Jittor团队的工作,期待其未来能够走得更好。


公众号近期荐读:

GAN整整6年了!是时候要来捋捋了! 

新手指南综述 | GAN模型太多,不知道选哪儿个?\

数百篇GAN论文已下载好!搭配一份生成对抗网络最新综述!

CVPR2020之MSG-GAN:简单有效的SOTA\

CVPR2020之姿势变换GAN:图像里谁都会劈叉?\

有点夸张、有点扭曲!速览这些GAN如何夸张漫画化人脸!\

见微知细之超分辨率GAN!附70多篇论文下载!\

天降斯雨,于我却无!GAN用于去雨如何?\

脸部转正!GAN能否让侧颜杀手、小猪佩奇真容无处遁形?\

容颜渐失!GAN来预测?\

强数据所难!SSL(半监督学习)结合GAN如何?\

弱水三千,只取你标!AL(主动学习)结合GAN如何?\

异常检测,GAN如何gan ?

虚拟换衣!速览这几篇最新论文咋做的!\

脸部妆容迁移!速览几篇用GAN来做的论文

【1】GAN在医学图像上的生成,今如何?

01-GAN公式简明原理之铁甲小宝篇


GAN&CV 交流群 ,无论小白还是大佬,诚挚邀您加入!\

一起讨论交流!长按备注【进群】加入:

更多分享、长按关注本公众号:\