清华「计图」现在支持国产芯片了!动态图推理比PyTorch快了270倍

新闻 前端
清华自研的深度学习框架计图(Jittor)在动态图推理速度上又一次完胜PyTorch。

本文经AI新媒体量子位(公众号ID:QbitAI)授权转载,转载请联系出处。

 清华自研的深度学习框架计图(Jittor)在动态图推理速度上又一次完胜PyTorch。

最近,计图团队完成了在寒武纪芯片MLU270上的移植。

这一次跟寒武纪的合作,使Jittor在backbone网络模型中的动态图推理速度较PyTorch平均提升了276.69倍

从团队公布的实验结果可以看到,在寒武纪芯片上分别用计图(Jittor)和PyTorch进行推理。

计图(Jittor)在16种backbone网络模型中的推理速度都较PyTorch大幅提升,其中包括alexnet、vgg系列、resnet系列。

其中最快的是alexnet,提升速度达到了464.43倍,最慢的resnet50也达到了153倍。

平均提升速度达276.69倍。

清华「计图」现在支持国产芯片了!动态图推理比PyTorch快了270倍

清华「计图」现在支持国产芯片了!动态图推理比PyTorch快了270倍

△表1:模型具体数值

与此同时,计图的精度损失也较PyTorch大幅减少。

清华「计图」现在支持国产芯片了!动态图推理比PyTorch快了270倍

△表2:resnet18的mse误差对比

团队以resnet18网络为例进行对比,baseline是CPU;结果显示:计图的精度损失比PyTorch降低了42.53%。

还是「元算子」的功劳

性能得到如此大的提升,是怎么实现的呢?

要知道,在TensorFlow,PyTorch框架底层,有上千个算子;想要移植新的硬件,就必须将大量的算子复写,这样一来不仅工作量剧增、而且难度也会加大。

Jittor这次还是依靠他们的核心武器——元算子

根据Jittor官方文档定义,元算子是指神经网络所需的基本算子。

早在设计Jittor的时候,研发团队就定下了一个目标,即用户只需要数行代码,就可定义新的算子和模型。在保证易用的同时,不丧失任何可定制性。

所以在Jittor中,多个元算子之间,可以相互融合成更加复杂的算子,这些复杂算子构成了神经网络计算的多个模块,如卷积层,归一化层等等。

清华「计图」现在支持国产芯片了!动态图推理比PyTorch快了270倍

△使用元算子实现卷积层

这一次将Jittor移植到寒武纪上,也是同样的原理。

研究团队把神经网络所需的基本算子,定义为三类共18个元算子。

让元算子相互融合形成常用算子,这样就能通过优化非常少的元算子,实现对常用算子性能的提升。

清华「计图」现在支持国产芯片了!动态图推理比PyTorch快了270倍

并且,元算子还是反向传播闭包,所有元算子反向传播后依旧是元算子。

所以在完成三类元算子的移植后,Jittor天然就能支持大部分常用元子的推理和简单训练。

从Jittor元算子到BANG

为了能让用户更为简单地操作,Jittor内置了元算子编译器,可以将用户的Python代码动态编译成寒武纪BANG语言

BANG语言是专门针对寒武纪产品架构的一种代码语言,它能极大优化寒武纪芯片的通用编程能力,提升用户编程的自由度。

并且它还有全套编译工具链来提高性能。包括CNCC(Cambricon Neuware Compiler Collection )、CNAS(Cambricon Neuware Assembler)、CNLINKER(Cambricon Neuware Linker)。

Jittor内置的元算子编译器可以把元算子自动转化为BANG算子。

采用动态编译的方式,能在运行时获取更多的信息,比如计算图上下文、形状信息等,这些信息都可以进一步提升算子的性能。

并且,Jittor还内置了可以根据硬件设备自动对BANG语言进行优化的优化编译遍(complier pass),从而生成对硬件设备更为适合的底层算子。

用这种方法,仅仅一行Python代码,就能表示BatchNorm算子的核心思想。

清华「计图」现在支持国产芯片了!动态图推理比PyTorch快了270倍

元算子编译器再把这行代码自动优化成BANG语言代码。

由于BANG语言的设计更加成熟,提供了类似于CUDA语言的线程调度模式,使得上手和调试都更加简单;也能更好地释放寒武纪芯片的算力。

清华「计图」现在支持国产芯片了!动态图推理比PyTorch快了270倍

清华自研,首个中国高校深度学习开源框架

计图(Jittor)的开发团队,均来自清华大学计算机系图形学实验室,负责人是清华大学计算机系的胡事民教授。

而主要负责开发的,则是来自实验室的博士生们:梁盾、杨国烨、杨国炜、周文洋……

计图(Jittor)与主流的深度学习框架TensorFlow、Pytorch等最大的不同在于,它是一个完全基于动态编译(Just-in-time)、使用元算子和统一计算图的深度学习框架。

它可以不像Pytorch那样依赖tracing机制,让用户的操作更加简单,同时还可以得到更好的学习训练效果。

[[398886]]

对于这一次取得的成果,开发者之一梁盾透露未来有可能发表在顶会上,并且会开源给大家。

此前,计图开源的点云模型库,在多种主流模型上训练性能较Pytorch提升一倍以上;并发布了第一个支持金属度、粗糙度的可微渲染库。

最近,他们还开源了一个智能P图神器DeepFaceEditing,可以通过草图自由编辑人脸。

清华「计图」现在支持国产芯片了!动态图推理比PyTorch快了270倍

感兴趣的同学可以去试试哦~

GitHub开源:https://github.com/IGLICT/DeepFaceEditing-Jittor

 

 

责任编辑:张燕妮 来源: 量子位
相关推荐

2020-12-09 14:06:04

计算机PyTorch 技术

2017-06-13 22:56:27

动态图PyTorchKeras

2019-03-15 11:34:19

神经网络数据图形

2021-05-13 15:23:31

人工智能深度学习

2017-12-13 10:08:26

大数据图数据推理数据科学

2018-04-11 09:46:12

存储芯片闪存

2021-07-06 06:26:43

动态计算图GPU深度学习

2020-09-17 17:46:20

Python地铁线路图

2023-10-25 18:53:45

芯片AI芯片

2020-06-02 10:09:59

JavaScript开发 可视化

2020-10-23 15:29:48

iPad抠图PC

2020-04-19 21:41:13

Python数据可视化

2023-12-11 15:40:32

PyTorch代码大模型

2021-06-18 15:50:56

模型人工智能深度学习

2018-06-14 16:59:42

TensorFlowEager深度学习

2017-02-08 17:08:32

谷歌深度学习计算图

2023-10-31 17:28:00

数据训练

2023-05-30 14:17:00

模型推理

2024-04-07 01:00:00

模型P图

2020-12-07 18:02:36

快手英特尔实时动态图训练效率
点赞
收藏

51CTO技术栈公众号