Jittor: 即时编译 (JIT) 的深度学习框架

快速开始 | 安装 | 教程 | 简体中文

Jittor 是一个基于 JIT 编译和元操作符的高性能深度学习框架。整个框架和元操作符都是即时编译的。Jittor 集成了强大的操作符编译器和调优器，使我们能够为您的模型生成高性能代码。Jittor 还包含丰富的高性能模型库，包括：图像识别、检测、分割、生成、可微渲染、几何学习、强化学习等。

前端语言是 Python。在前端使用模块设计和动态图执行，这是深度学习框架接口中最受欢迎的设计。后端由高性能语言实现，如 CUDA、C++。

相关链接：

• Jittor 网站
• Jittor 教程
• Jittor 模型
• Jittor 文档
• Github, GitLink, Gitee
• Jittor 论坛
• Awesome Jittor 清单
• IM: QQ 群(761222083)

以下示例展示了如何一步一步地建模一个两层神经网络并从头开始训练，只需几行 Python 代码。

import jittor as jt
from jittor import Module
from jittor import nn
import numpy as np

class Model(Module):
    def __init__(self):
        self.layer1 = nn.Linear(1, 10)
        self.relu = nn.Relu() 
        self.layer2 = nn.Linear(10, 1)
    def execute (self,x) :
        x = self.layer1(x)
        x = self.relu(x)
        x = self.layer2(x)
        return x

def get_data(n): # 生成用于训练测试的随机数据
    for i in range(n):
        x = np.random.rand(batch_size, 1)
        y = x*x
        yield jt.float32(x), jt.float32(y)

learning_rate = 0.1
batch_size = 50
n = 1000

model = Model()
optim = nn.SGD(model.parameters(), learning_rate)

for i,(x,y) in enumerate(get_data(n)):
    pred_y = model(x)
    dy = pred_y - y
    loss = dy * dy
    loss_mean = loss.mean()
    optim.step(loss_mean)
    print(f"step {i}, loss = {loss_mean.data.sum()}")

内容

• 快速开始
• 安装
• 教程
• 贡献
• 团队
• 许可

快速开始

我们提供了一些 jupyter 笔记本以帮助您快速开始使用 Jittor。

• 示例：模型定义和训练
• 基本知识：操作符，变量
• 元操作符：使用元操作符实现自己的卷积

安装

Jittor 环境要求：

Jittor 提供三种安装方式：pip、docker 或手动。

Pip 安装

sudo apt install python3.7-dev libomp-dev
python3.7 -m pip install jittor
# 或从 github 安装最新版本
# python3.7 -m pip install git+https://github.com/Jittor/jittor.git
python3.7 -m jittor.test.test_example

macOS 安装

请首先使用 homebrew 安装额外的依赖项。

brew install libomp

然后您可以通过 pip 安装 jittor 并运行示例。

python3.7 -m pip install jittor
python3.7 -m jittor.test.test_example

目前 jittor 仅支持在 macOS 上使用 CPU。

Windows 安装

# 检查您的 Python 版本 (>=3.8)
python --version
python -m pip install jittor
# 如果使用 conda
conda install pywin32

在 Windows 上，jittor 会自动检测并安装 CUDA，请确保您的 NVIDIA 驱动程序支持 CUDA 10.2 或更高版本，或者您可以手动让 jittor 为您安装 CUDA：

python -m jittor_utils.install_cuda

Docker 安装

我们提供了一种 Docker 安装方法以避免您配置环境。Docker 安装方法如下：

# 仅 CPU (Linux)
docker run -it --network host jittor/jittor
# CPU 和 CUDA (Linux)
docker run -it --network host --gpus all jittor/jittor-cuda
# 仅 CPU (Mac 和 Windows)
docker run -it -p 8888:8888 jittor/jittor

手动安装

我们将展示如何在 Ubuntu 16.04 上一步步安装 Jittor，其他 Linux 发行版可能有类似的命令。

第一步：选择您的后端编译器

# g++
sudo apt install g++ build-essential libomp-dev

# 或 clang++-8
wget -O - https://raw.githubusercontent.com/Jittor/jittor/master/script/install_llvm.sh > /tmp/llvm.sh
bash /tmp/llvm.sh 8

第二步：安装 Python 和 python-dev

Jittor 需要 Python 版本 >= 3.7。

sudo apt install python3.7 python3.7-dev

步骤3：运行Jittor

整个框架是即时编译的。让我们通过pip安装jittor

git clone https://github.com/Jittor/jittor.git
sudo pip3.7 install ./jittor
export cc_path="clang++-8"
# 如果使用其他编译器，则更改cc_path
# export cc_path="g++"
# export cc_path="icc"

# 运行一个简单的测试
python3.7 -m jittor.test.test_example

如果测试通过，说明你的Jittor已经准备好了。

可选步骤4：启用CUDA

在Jittor中使用CUDA非常简单，只需设置环境变量nvcc_path

# 将此变量替换为你的nvcc位置
export nvcc_path="/usr/local/cuda/bin/nvcc" 
# 运行一个简单的cuda测试
python3.7 -m jittor.test.test_cuda 

如果测试通过，你可以通过设置use_cuda标志来在Jittor中使用CUDA。

import jittor as jt
jt.flags.use_cuda = 1

可选步骤5：测试Resnet18训练

为了检查Jittor的完整性，你可以运行Resnet18训练测试。注意：此测试需要6G GPU RAM。

python3.7 -m jittor.test.test_resnet

如果这些测试未通过，请向我们报告错误，并随时贡献代码 ^_^

教程

在教程部分，我们将简要解释Jittor的基本概念。

要使用Jittor训练你的模型，你只需了解三个主要概念：

• Var：jittor的基本数据类型
• 操作：Jittor的操作与numpy相似

Var

首先，让我们从Var开始。Var是jittor的基本数据类型。Jittor中的计算过程是异步的，以进行优化。如果你想访问数据，可以使用Var.data进行同步数据访问。

import jittor as jt
a = jt.float32([1,2,3])
print (a)
print (a.data)
# 输出：float32[3,]
# 输出：[ 1. 2. 3.]

我们还可以给变量命名。

a.name('a')
print(a.name())
# 输出：a

操作

Jittor的操作与numpy相似。让我们尝试一些操作。我们通过操作jt.float32创建了Var a和b，并将它们相加。打印这些变量显示它们具有相同的形状和数据类型。

import jittor as jt
a = jt.float32([1,2,3])
b = jt.float32([4,5,6])
c = a*b
print(a,b,c)
print(type(a), type(b), type(c))
# 输出：float32[3,] float32[3,] float32[3,]
# 输出：<class 'jittor_core.Var'> <class 'jittor_core.Var'> <class 'jittor_core.Var'>

此外，我们使用的所有操作jt.xxx(Var, ...)都有别名Var.xxx(...)。例如：

c.max() # jt.max(c)的别名
c.add(a) # jt.add(c, a)的别名
c.min(keepdims=True) # jt.min(c, keepdims=True)的别名

如果你想知道Jittor支持的所有操作，请尝试help(jt.ops)。所有你在jt.ops.xxx中找到的操作，都可以通过别名jt.xxx使用。

help(jt.ops)
# 输出：
#   abs(x: core.Var) -> core.Var
#   add(x: core.Var, y: core.Var) -> core.Var
#   array(data: array) -> core.Var
#   binary(x: core.Var, y: core.Var, op: str) -> core.Var
#   ......

更多

如果你想了解更多关于Jittor的信息，请查看以下笔记本：

• 快速入门
  • 示例：模型定义和训练
  • 基础知识：操作, Var
  • 元操作：使用元操作器实现你自己的卷积
• 高级
  • 自定义操作：用C++和CUDA编写你的操作并即时编译
  • 性能分析器：分析你的模型
  • Jtune: 性能调优工具

你可以在自己的电脑上通过python3.7 -m jittor.notebook启动这些笔记本。

贡献

Jittor还很年轻，可能包含错误和问题。请在我们的错误跟踪系统中报告它们。欢迎贡献。此外，如果你对Jittor有任何想法，请告诉我们。

你可以通过以下方式帮助Jittor：

• 在你的论文中引用Jittor
• 向你的朋友推荐Jittor
• 贡献代码
• 贡献教程和文档
• 提交问题
• 回答jittor相关问题
• 点亮星星
• 关注jittor
• ......

联系我们

网站: http://cg.cs.tsinghua.edu.cn/jittor/

邮箱: jittor@qq.com

提交问题: https://github.com/Jittor/jittor/issues

QQ群: 836860279

团队

Jittor目前由Tsinghua CSCG Group维护。如果你对Jittor感兴趣并想改进它，请加入我们！

引用

@article{hu2020jittor,
  title={Jittor: a novel deep learning framework with meta-operators and unified graph execution},
  author={Hu, Shi-Min and Liang, Dun and Yang, Guo-Ye and Yang, Guo-Wei and Zhou, Wen-Yang},
  journal={Science China Information Sciences},
  volume={63},
  number={222103},
  pages={1--21},
  year={2020}
}

许可证

Jittor按照LICENSE.txt文件中的Apache 2.0许可进行许可。