Jittor/jittor

Jittor – 一種即時編譯的深度學習框架

是什麼 – Jittor 是一個以 C++/CUDA 寫成的全棧深度學習庫，並提供 Python 前端。它將整個計算圖 即時編譯（JIT）並使用元操作符系統，為您正在訓練的特定模型生成高度優化的內核。設計遵循 PyTorch/TensorFlow-Eager 中熟悉的動態圖風格，您只需撰寫一般的 Python 程式碼，框架就會自動將其轉換為高效能的原生程式碼。

核心理念

• 即時編譯 – 每個運算子及其周圍圖皆在執行時即時編譯，讓編譯器能針對精確的張量形狀與硬體生成特化程式碼。
• 元操作符 – 高階的建構模組，可組合或自訂，讓使用者不需離開 Python 即可撰寫新運算子（例如自訂卷積）。
• 統一圖執行 – Jittor 建構動態圖，但也能融合運算以達成最佳效能，類似於靜態圖框架。
• 後端彈性 – 支援 CPU、NVIDIA CUDA、AMD ROCm 以及 Hygon DCU；核心為 C++/CUDA，使用者 API 為純 Python。

可以做什麼

• 直接訓練經典視覺模型（ResNet、檢測、分割、生成模型），無需額外設定。
• 使用提供的模型庫，實驗可微分渲染、幾何深度學習、強化學習等。
• 以 C++/CUDA 寫出高效率的自訂運算子，並自動進行 JIT 編譯。
• 使用內建的 jtune 工具進行效能分析與調校。

安裝

• Pip – python -m pip install jittor（Linux 需要 python3-dev 與 libomp-dev）。
• Docker – 預先建置的映像檔 jittor/jittor（僅 CPU）與 jittor/jittor-cuda（GPU）。
• 手動 – 克隆程式碼庫，安裝 C++ 編譯器（g++ ≥ 5.4 或 clang ≥ 8），安裝 Python ≥ 3.7，然後 sudo pip install ./jittor。
• GPU 支援為可選；安裝 CUDA 後設定 jt.flags.use_cuda = 1，或在 Windows 上讓 Jittor 自動下載。

快速範例（兩層迴歸模型） – README 包含一個最小腳本，展示典型工作流程：

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

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

model = Model()
opt = nn.SGD(model.parameters(), lr=0.1)

for x, y in data_generator():
    pred = model(x)
    loss = ((pred - y) ** 2).mean()
    opt.step(loss)
    print('loss', loss.data.item())

程式碼與 PyTorch 風格模組相似，但底層 Jittor 會 JIT 編譯整個前向-反向傳播以達成最大速度。

學習資源

• 官方網站與文件： https://cg.cs.tsinghua.edu.cn/jittor/
• 用於基礎、元操作符、自訂運算子、效能分析等的互動式筆記本。
• 社群論壇、QQ 群，以及第三方專案「Awesome Jittor」清單。

生態系統

• 模型庫涵蓋影像分類、檢測、分割、生成模型、可微分渲染、幾何學習與強化學習。
• 工具：jtune 用於內核調校，效能分析器，以及模型匯出/匯入工具。
• 可透過自訂 C++/CUDA 運算子與 Python 元操作符擴充。

授權 – Apache 2.0（參見 LICENSE.txt 檔案）。

引用 – 若在研究中使用 Jittor，請引用 README 中提供的 2020 年《中國科學：資訊科學》論文。