领域相关,也来凑个热闹。
看到这个工作还是比较开心的,关于AI技术,国内以及华人圈子已经在发起越来越多具备硬核属性的努力了。
硬件上,有海思的Da Vinci,燧原的IPU,中科院的寒武纪(以及我司的含光)等等,软件上,有TVM,头条&港大合作的BytePS,商汤的性能优化团队(除了不开源都挺好)的底层优化工作(说手工优化似乎不完全合适,还是说手工优化+半手工优化+AI编译优化比较合适^-^),华为在AI底层软件方面的工作(在AI编译器方面,华为是国内最早投入的公司之一),以及这次的Jittor的工作(也不要落了进辉兄的Oneflow工作,能够在AI Infra这种硬核方向坚持创业多年,是很值得尊敬的努力),当然内举不避亲,我们阿里PAI团队在AI硬核方向也投入不少,分布式、资源优化、编译优化、模型压缩、强化学习框架、迁移学习平台、建模辅助工具等等,欢迎感兴趣的同学勾搭,我可以帮忙route到合适的团队里去(此处满满的招聘带货嫌疑^-^)。
回归正题。
作为一个探索性的工作,我觉得Jittor是值得尊敬的,能够在这个工作里看到不少好的创新点。因为我对底层关注相对多一些,所以直接举一些偏底层的features,包括:
统一内存我们去年在生产集群里出于集群资源优化的需要,也做了类似的feature(@孙敏敏 ),也在训练作业里,发现有一些用户比较喜欢这个feature(麻麻我的作业不用再担心显存不够用被钉,只是可能会变慢^-^)。
跨迭代融合在目前我所知道的AI框架里还不支持(当然可以通过对现有框架hack来达到这个效果),而跨迭代融合功能的缺失会导致一些inter-minibatch的优化机会的丢失,这一点我没记错的话,去年在BytePS的工作里似乎还专门提到过(P3里是不是有提到映象不深,但本质上是相关的)。
新AI编译框架的尝试,倒不用多说,看到code和sample就足够self-explaining了。能够自己定义一套DSL,基于这套DSL来描述算子生成的空间,并打通完整的通路,也具备了一定的fusion能力的支持,即便只是在一些示例模型上跑通,这个工作也还是值得尊敬和认可的。
总的来说,值得赞赏。接下来我也会表达一些其他的comments:
最后,还是很appreciate Jittor团队的工作,期待其未来能够走得更好。
也欢迎来勾搭/交流/合作,多多响应扬清回复里的号召:)。
利益相关:Jittor开发者梁盾七年多的室友。
我在清华计算机系是个菜鸡,有大佬大神在轮不到我来评价Jittor的优缺点,这个答案只是想安利各位同行尝试一下这个框架,多给开发团队提意见。盾神是我认识的人里最有极客范儿的之一,热爱钻研自己感兴趣的技术并且享受实现的过程,不追热点不水论文不出风头,关注CS领域的最新技术然后去学习去看文档去看源码去实现,最后通过自己的努力做出真正实质性的贡献去改变,这整个过程就是他最大的乐趣来源。我一直认为这才是真正优秀的计算机系学生,既佩服又羡慕他。盾神高中就玩斯坦福兔子,大学拿超算世界冠军,读博又深入底层做了这样优秀的深度学习框架,但为人还是一如既往地低调。我以为盾神会用自己的知乎账号答这个题,没想到他们注册了一个叫Jittor的账号。对于这样的人,最开心的事莫过于自己辛辛苦苦做出来的东西能够被更多的人使用、能够有自己的使用者社区。而且盾神是我认识的同级清华计算机系博士生里最强的,我相信他花这么多时间精力出来的成果也不会让人失望,所以真的强烈推荐大家去试着用一下,毕竟从pytorch上switch过去也很方便。
盾神开发Jittor的这一年多是我见过的他工作最辛苦的一段时间,作为一个在清华计算机系被称为神的家伙以及前OI国家队候补队员,盾神和其他的神一样,对于我们系专业课的作业都是随便花点时间就拿高分的,而且盾神代码和工程能力很强,基本上看不到他有需要每天码很久连续码很多天的情况。在写Jittor之前,盾神在CV和机器人方面做过一些工作,对tf和pytorch钻研颇多,而且跟我们这些调包怪不同,他对底层乃至硬件上的机制也都足够了解。但是Jittor不是一个大作业也不是一篇随便水水的论文,在AI会议动辄几千上万篇投稿的今天,想要做出真正有用的东西需要付出巨大的努力,一个重新设计的深度学习框架需要学很多东西写很多代码,对于一个学生团队来说非常不容易。这一年多看盾神从C++写到binding写到汇编写到算子库Cache硬件各种优化再写到接口,在这个过程中甚至还顺便搞了个公众号写优化方面的blog。Jittor经过反复迭代测试今天终于发布,作为好朋友很为他高兴。
盾神迭代开发测试过程中还设计了问卷,调查开发者们对深度学习框架的需求和使用习惯,之所以Jittor在接口上跟pytorch如此相似也是因为大部分开发者更习惯这样。没错,我说的唯一一句话是填这个问卷超出了我的能力范围,可见我有多菜。但即使是这么菜的我,盾神也找我聊了很多次下一步框架应该怎么扩展,大家可能还有哪些需求,他是真的希望花了这么多时间精力做出来的东西能够实实在在地有用,所以我相信他们团队一定会重视每一条有意义的issue,后续也会推出更多好用的模型和接口,扩展Jittor的适用领域,再次推荐大家多尝试多关注。
计图(Jittor),清华“贵系”(计算机系)图形实验室出品,是首个来自中国高校科研机构的开源深度学习框架。此前,业内来自“高校”的框架还有加拿大蒙特利尔大学的Theano,UC伯克利的Caffe。
与主流的深度学习框架TensorFlow、Pytorch不同,Jittor是一个完全基于动态编译(Just-in-time)、使用元算子和统一计算图的深度学习框架。
研发团队介绍称,开发Jittor是为了将新技术、硬件和模型的能力,更好地释放出来。
“深度学习发展迅猛,TensorFlow、PyTorch这些老牌主流框架,也会在新模型,新算法,新硬件上表现不佳,所以需要新的框架,在易于扩展同时保持高效。”
而现在框架呈现出来的能力,的确有超越“前辈”的倾向:
基于Jittor开发的深度学习模型,可以实时自动优化并运行在指定的硬件上,如CPU,GPU,在多种机器视觉任务上能够比同类产品PyTorch性能提高10%~50%。
团队还介绍,如此成果,得益于Jittor的两大创新点:元算子和统一计算图。这也是Jittor的立身之本。
根据Jittor官方文档定义,元算子是指神经网络所需的基本算子。
在TensorFlow,PyTorch框架底层,有上千个算子,如此多的算子使得开发和优化难度大幅提升。
在设计Jittor的时候,他们就定下了一个目标,即用户只需要数行代码,就可定义新的算子和模型。同时在保证易用的同时,不丧失任何可定制性。
所以在Jittor中,多个元算子之间,可以相互融合成更加复杂的算子,这些复杂算子构成了神经网络计算的多个模块,如卷积层,归一化层等等。
他们将这种融合称为元算子融合,可以提升性能,节省资源。在文档中,他们分享了一个案例:只用4个元算子,就实现了卷积操作。
他们介绍称,元算子的可拓展性很强,通过对元算子的简单修改,就可以实现更多复杂的卷积操作,如扩张卷积、深度卷积、点卷积、分离式卷积、反卷积等。
而且,通过元算子反向传播闭包,能自动生成反向卷积层。具体如下图所示,反向卷积层将来自输出的梯度,通过4个元算子,将梯度反向传播给卷积层的权重:
Jittor开发团队介绍称,在这样的设计下,元算子和Numpy一样易于使用,并且超越Numpy能够实现更复杂更高效的操作。
而且,通过元算子的反向传播闭包,Jittor可以对所有前向反向算子进行统一管理,这就是他们所说的第二个创新点:统一计算图。
简单来说,统一计算图是完成了多种统一的动态计算图。根据官方文档介绍,在Jittor中,核心有四个方面的统一:
基于这个方面,他们给出了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的整体架构一共分为四层,如下图所示:
它是基于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%。
此外,为了方便更多人上手Jittor,开发团队采用了和PyTorch较为相似的模块化接口,并提供辅助转换脚本,可以将PyTorch的模型自动转换成Jittor的模型。
他们介绍称,在参数保存和数据传输上,Jittor使用和PyTorch一样的 Numpy+pickle 协议,所以Jittor和PyTorch的模型可以相互加载和调用。
当然, Jittor作为一个新兴深度学习框架,在一些功能上,仍旧需要持续迭代完善。比如生态的建设,以及更大范围的推广,仍旧需要很多的努力。
Jittor开发团队介绍称,就目前来看,Jittor框架的模型支持还待完善,分布式功能待完善。这也是他们下一阶段研发的重点。
最后,是时候介绍Jittor的开发团队出场,他们来自清华大学计算机系图形学实验室,牵头者是清华大学计算机系胡事民教授。
该实验室的主要研究方向是计算机图形学、计算机视觉、智能信息处理、智能机器人、系统软件等,在ACM TOG, IEEE TVCG, IEEE PAMI, ACM SIGGRAPH, IEEE CVPR, IEEE ICRA, USENIX ATC等重要国际刊物上发表论文100余篇。
开发Jittor的主力,是该实验室梁盾、杨国烨、杨国炜、周文洋等一批博士生。
据梁盾透露,他们接下来的计划,是先围绕学界,重点发力。希望能成为国内以及世界上学术界最受欢迎,使用最多的框架,并对AI产业界产生积极的影响。
但想要走通这条路,并没有那么容易。TensorFlow和PyTorch已经成为了当前主流的深度学习框架,正在被全世界的研究者们采用。尤其是PyTorch,正在大面积抢占学术界。
其实从模型特性,以及设计理念来看,PyTorch可能是Jittor更直接的对标对象。
Jittor将如何发力?
在他们的规划中,接下来将组建开源社区,除了完善框架外,还会联合多所高校使用Jittor教授课程,以现有人员作为核心,壮大开发团队和用户,首要目标是服务更多研究人员。
据说,已经有多位高校教授,决定要在自己课堂上使用。
同时,另一个公开信息也值得关注:胡事民教授从2010年开始,就担任清华大学—腾讯联合实验室主任。在Jittor研发过程中,还得到了这一实验室支持。
所以这一框架是否会与腾讯展开合作?
目前研究团队没有给出直接明确的答复,但表示:非常希望能和更多的产业界的伙伴们联手推动Jittor的发展。
总之,打造AI框架本身已不易,开源之后更要接受各方直接检验。
现在,清华迈出了关键一步,虚的不多说,各位收好下方传送门,走过路过不要错过,都参与检验一下吧~
https:// cg.cs.tsinghua.edu.cn/j ittor/
—完—
@量子位 · 追踪AI技术和产品新动态
深有感触的朋友,欢迎赞同、关注、分享三连վ'ᴗ' ի ❤
感谢各位专家的评价~不胜荣幸
刚刚发布收到了很多小伙伴的关注和疑问,我们研发团队也补充一下。
我先来说一下背景,现在深度学习框架主流的框架包括tf,pytorch,国内现在是深度学习研究的主力,但是国内急需更好的深度学习的生态,希望我们开源的深度学习框架能像星星之火可以燎原,也希望能够和学界和产业界共建和谐生态。
从我们自己官方网站 Jittor: 即时编译深度学习框架 中把设计理念抄过来:
易用且可定制 用户只需要数行代码,就可定义新的算子和模型,在易用的同时,不丧失任何可定制性。
实现与优化分离 用户可以通过前端接口专注于实现,而实现自动被后端优化。从而提升前端代码的可读性,以及后端优化的鲁棒性和可重用性。
所有都是即时的 Jittor的所有代码都是即时编译并且运行的,包括Jittor本身。用户可以随时对Jittor的所有代码进行修改,并且动态运行。
期望Jittor能为学界和业界提供一个灵活高效的深度学习平台,促进人工智能的研究和应用,赋能人工智能产业。
同时,Jittor才刚刚开始,前面的路还很长,平台的开发需要对图形图像应用的精准把握、对机器学习算法的深刻理解、对底层系统软件的巧妙应用。Jittor可能还存在一些错误和问题,并需要不断完善和演进,我们会一直努力,并欢迎批评和各种建议。
感谢大家对Jittor的关注和评论,对于很多小伙伴提到关于学习成本的问题,Jittor在设计时其实也做了考虑,并进行了用户调研,最终决定使用与pytorch类似的接口,Jittor在方法名称和参数的设计时也尽量和pytorch保持一致。我们已经提供了API文档,对于简单的应用,我们相信有机器学习经验的小伙伴,能够轻松上手jittor!
欢迎机器学习的爱好者加入我们。
了解更多信息,请访问“计图(Jittor)”
欢迎大家去github提交issue还有讨论:Jittor/jittor
官方网站。https://cg.cs.tsinghua.edu.cn/jittor。
Email:jittor@qq.com。
QQ群:761222083
最新动态:
挺多朋友给我发这个项目,我粗略研究了一下,简单单理解是:
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 补充一些关于前端的思考
简单看了一下代码,非常有意思的一个项目。因为看得不深入,所以如果有观察错误的话,我想对作者先致以歉意。
总的来说:项目更加关注在如何进行计算图优化以及just in time compilation上面(所以叫jittor),并不是关注完整的端到端的框架设计(比如说建模前端等等),我觉得定位是比较清楚的,自动代码生成是现在大家都很关注的方向,在校的同学能够着手把这一套都做一次,值得点赞。
一些能看到的工程点:
很赞的一些工程细节:
最后再次表示:非常有意思的一个项目,可以看到作者在学校关注最近的科研方向,同时又能很hands on的自己造一把轮子(褒义)的能力。
打个广告(团队同学已经比我手快了):欢迎对机器学习、大数据、分布式系统等技术有兴趣的同学来一起实习、工作,或者在开源社区合作,夏季实习请扫码 :)