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 논문을 인용해 주세요.