计图(Jittor)是一个完全基于动态编译(Just-in-time),内部使用创新的元算子和统一计算图的深度学习框架。元算子和 Numpy 一样易于使用,而统一计算图则是融合了静态计算图和动态计算图的诸多优点,在易于使用的同时,提供高性能的优化。基于元算子开发的深度学习模型,可以被计图实时的自动优化并且运行在指定的硬件上,如 CPU、GPU。

特性:

  • 算子动态编译:Jittor 内置元算子编译器,可以将用户通过元算子编写的 Python 代码,动态编译成高性能的C++ 代码。
  • 自动优化:jittor 内置优化编译遍(complier pass),同时和 LLVM 兼容,这些编译遍会根据硬件设备自动优化动态编译的代码,常见的优化编译遍有循环重排、循环分裂、循环融合、数据打包、向量化、GPU并行。这些编译遍对 C++ 代码进一步优化,生成对计算设备友好的底层算子。
  • 统一内存管理:Jittor 使用了统一内存管理,统一 GPU 与 CPU 之间的内存,当深度学习模型将 GPU 内存耗尽时,将使用 CPU 内存来弥补。
  • 高效同步异步接口:Jittor 同时提供两种接口,同步和异步接口之间切换不会产生性能损失。
  • 模型迁移:Jittor 采用和 PyTorch 相似的模块化接口,同时提供辅助转换脚本,可将 PyTorch 的代码转换成 Jittor 的模型,并可以相互加载和调用。
  • 元算子融合:可将神经网络的基本算子(元算子)融合成复杂算子,如卷积层、归一化层,进一步构成神经网络和深度学习应用。
  • 高阶导数以及反向传播闭包:元算子是反向传播闭包(元算子反向传播也是元算子),同时支持计算任意高阶导数。在深度学习算子开发过程中,免去反向传播算子重复开发工作,同时可以使用统一的优化策略。

安装:

Jittor前端语言为Python,使用了模块化的设计,类似于PyTorch,Keras;后端则使用高性能语言编写,如CUDA,C++。当前Jittor支持三种即时编译器:

  • CPU 编译器 (需要下列至少一个)
    • g++ (>=5.4.0)
    • clang (>=8.0)推荐
  • GPU 编译器(可选)
    • nvcc(>=10.0)