近日，计图(Jittor)团队成功完成寒武纪芯片MLU上的移植，支持推理和训练，并复现了ResNet、Alexnet、VGG等骨干网络，其中动态图推理速度相较PyTorch平均提升了.69倍，推理精度也得到了显著的提升。Jittor是清华大学自主研发的框架，相比PyTorch具有更好的可移植性，在保证易用的同时，能更充分发挥寒武纪芯片的性能，此次合作有望革新现有软硬件体系架构，进一步突破AI算力瓶颈，加速AI研究的创新创造，为国内的人工智能生态打下坚实基础。

Part1

多个骨干网络速度提升数百倍Jittor团队在寒武纪芯片上与PyTorch对比了多种backbone网络模型的速度，包括alexnet、vgg系列、resnet系列在内的16个网络。

可视化结果如图1所示，动态图性能(逐层)，Jittor的速度可以达到PyTorch的~倍，平均速度达到了PyTorch的.69倍。其中最快的是alexnet，性能提升可达.43倍。

图1不同backbone网络Jittor动态图加速比

每个模型的具体数值如表1所示，其中Model一列展示了不同backbone模型，测试使用长宽均为的rgb图像作为输入。Jittor和PyTorch两列分别展示了两个框架单次前向所需的时间，单位是毫秒。SpeedUp展示了Jittor动态图的加速比。

表1Jittor和PyTorch在不同backbone的动态图速度指标同时发现Jittor的精度损失也会比PyTorch大幅减少。以Resent18网络为例，对比了Jittor和PyTorch的MSE，结果如表2所示。对比的baseline是cpu，测试表明，Jittor的精度损失比PyTorch少42.53%。表2resnet18的mse误差对比

Part2

Jittor如何助力寒武纪性能提升？

主流的TensorFlow和PyTorch框架有几百甚至上千个算子，这就导致了在移植新的硬件时需要付出巨大的工作量一一复现这些算子。

而Jittor采用元算子的的概念，将神经网络所需的基本算子定义为三类共18个元算子，这些元算子能相互融合成大部分常用的算子[1]。因此对少量元算子进行优化，就能使得不同的常用算子性能都得到显著提升。

除此之外，元算子还是反向传播闭包，这意味着所有元算子的反向传播算子仍然是元算子，也就是当完成移植元算子的移植后，Jittor就天然支持训练。

因此，得益于元算子设计，在完成三类元算子的移植后，Jittor就支持了大部分常用算子的推理和简单训练(见图2)。

图2在不同硬件移植元算子以支持推理和训练Jittor的多个元算子之间可以相互融合成更加复杂的算子，这些复杂算子构成了神经网络计算的多个模块，如卷积层，归一化层等(见图3)。这种融合我们称之为元算子融合，可以提升性能，节省资源。

传统的算子融合方式需要手工编写融合规则，并逐一开发融合后的算子。而Jittor通过实时分析计算图结构，自动将可以融合的元算子进行融合，从而大幅减少访存带来的开销，并减少开发工作量。

图3使用元算子实现卷积层的方法示意图Jittor内置了元算子编译器，可以通过动态编译的方式将用户的Python代码编译成高性能的寒武纪BANG代码。

Jittor的动态编译过程相比较传统的静态编译，可以在运行时获得更多的额外信息，如计算图上下文，形状信息等等，这些信息都可以进一步用于提升算子性能。

此外，Jittor还内置了优化编译遍(