清华「计图」现在支持国产芯片了！动态图推理比PyTorch快了270倍
清华「计图」现在支持国产芯片了！动态图推理比PyTorch快了270倍
明敏发自凹非寺
量子位报道公众号 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 倍。
△表1：模型具体数值
与此同时，计图的精度损失也较 PyTorch 大幅减少。
△表2：resnet18 的 mse 误差对比
团队以 resnet18 网络为例进行对比，baseline 是 CPU；结果显示：计图的精度损失比 PyTorch 降低了 42.53%。
还是「元算子」的功劳
性能得到如此大的提升，是怎么实现的呢？
要知道，在 TensorFlow，PyTorch 框架底层，有上千个算子；想要移植新的硬件，就必须将大量的算子复写，这样一来不仅工作量剧增、而且难度也会加大。
Jittor 这次还是依靠他们的核心武器——元算子。
根据 Jittor 官方文档定义，元算子是指神经网络所需的基本算子。
早在设计 Jittor 的时候，研发团队就定下了一个目标，即用户只需要数行代码，就可定义新的算子和模型。在保证易用的同时，不丧失任何可定制性。
所以在 Jittor 中，多个元算子之间，可以相互融合成更加复杂的算子，这些复杂算子构成了神经网络计算的多个模块，如卷积层，归一化层等等。
△使用元算子实现卷积层
这一次将 Jittor 移植到寒武纪上，也是同样的原理。
研究团队把神经网络所需的基本算子，定义为三类共 18 个元算子。
让元算子相互融合形成常用算子，这样就能通过优化非常少的元算子，实现对常用算子性能的提升。
并且，元算子还是反向传播闭包，所有元算子反向传播后依旧是元算子。
所以在完成三类元算子的移植后，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 算子的核心思想。
元算子编译器再把这行代码自动优化成 BANG 语言代码。
由于 BANG 语言的设计更加成熟，提供了类似于 CUDA 语言的线程调度模式，使得上手和调试都更加简单；也能更好地释放寒武纪芯片的算力。
清华自研，首个中国高校深度学习开源框架
计图（Jittor）的开发团队，均来自清华大学计算机系图形学实验室，负责人是清华大学计算机系的胡事民教授。
而主要负责开发的，则是来自实验室的博士生们：梁盾、杨国烨、杨国炜、周文洋……
计图（Jittor）与主流的深度学习框架 TensorFlow、Pytorch 等最大的不同在于，它是一个完全基于动态编译（Just-in-time）、使用元算子和统一计算图的深度学习框架。
它可以不像 Pytorch 那样依赖 tracing 机制，让用户的操作更加简单，同时还可以得到更好的学习训练效果。
对于这一次取得的成果，开发者之一梁盾透露未来有可能发表在顶会上，并且会开源给大家。
此前，计图开源的点云模型库，在多种主流模型上训练性能较 Pytorch 提升一倍以上；并发布了第一个支持金属度、粗糙度的可微渲染库。
最近，他们还开源了一个智能P图神器 DeepFaceEditing，可以通过草图自由编辑人脸。
感兴趣的同学可以去试试哦~
来自: page.om.qq.com