Jittor/jittor

Jittor – ジャストインタイムでコンパイルされるディープラーニングフレームワーク

何であるか – Jittor は C++/CUDA で書かれたフルスタックのディープラーニングライブラリであり、Python フロントエンドを備えています。計算グラフ全体を ジャストインタイム（JIT）でコンパイルし、特定のモデルに最適化されたカーネルを生成するためにメタ演算子システムを使用します。PyTorch/TensorFlow-Eager で使われる馴染み深い動的グラフスタイルに従っているため、通常の Python コードを書くだけで、フレームワークが高速なネイティブコードに変換してくれます。

主なアイデア

• JIT コンパイル – 各演算子と周囲のグラフが実行時にコンパイルされ、コンパイラが正確なテンソル形状とハードウェアに特化したコードを生成できます。
• メタ演算子 – 高レベルな構成ブロックで、組み合わせたりカスタマイズしたりでき、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 が自動的にダウンロードします。

クイック例（2層回帰モデル） – 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」リスト。

エコシステム

• 画像分類、検出、セグメンテーション、生成モデル、微分可能レンダリング、幾何学的学習、RL をカバーするモデルズー。
• ツール：jtune（カーネルチューニング）、プロファイラー、モデルのエクスポート/インポートユーティリティ。
• カスタム C++/CUDA 演算子と Python メタ演算子による拡張が可能。

ライセンス – Apache 2.0（LICENSE.txt ファイルを参照）。

引用 – 研究で Jittor を使用する場合、README に記載されている 2020 年の Science China Information Sciences 論文を引用してください。