tinygrad MNIST 튜토리얼 — 분류 모델 훈련하기

tinygrad MNIST 튜토리얼 — 분류 모델 훈련하기

tinygrad를 설치한 뒤의 첫 튜토리얼로 딱 좋아요. 노트북이나 colab에서 시작하면 됩니다. tinygrad는 아주 가벼워 어디서든 쉽게 설치되고 특별한 colab 이미지가 필요 없지만, 속도를 위해 T4 GPU 이미지를 권장해요.

출처: https://docs.tinygrad.org/mnist

설치

colab에서의 한 줄 설치는:

! pip install git+https://github.com/tinygrad/tinygrad.git

기본 디바이스는 이렇게 확인해요. GPU 인스턴스에서는 CUDA, CPU 인스턴스에서는 CPU가 보입니다.

from tinygrad import Device
print(Device.DEFAULT)

간단한 모델

Keras 튜토리얼의 모델을 가져와 사용해요. PyTorch와 두 가지 큰 차이가 있습니다. stateful 레이어만 __init__에 선언 한다는 점과, nn.Module 클래스나 forward 함수가 아니라 그냥 일반 클래스와 __call__ 을 쓴다는 점입니다.

from tinygrad import Tensor, nn, Context

class Model:
    def __init__(self):
        self.l1 = nn.Conv2d(1, 32, kernel_size=(3, 3))
        self.l2 = nn.Conv2d(32, 64, kernel_size=(3, 3))
        self.l3 = nn.Linear(1600, 10)

    def __call__(self, x: Tensor) -> Tensor:
        x = self.l1(x).relu().max_pool2d((2, 2))
        x = self.l2(x).relu().max_pool2d((2, 2))
        return self.l3(x.flatten(1).dropout(0.5))

데이터셋 가져오기

tinygrad는 MNIST를 내장하고 있어서 네 줄이면 끝나요.

from tinygrad.nn.datasets import mnist

X_train, Y_train, X_test, Y_test = mnist()
print(X_train.shape, X_train.dtype, Y_train.shape, Y_train.dtype)
# (60000, 1, 28, 28) dtypes.uchar (60000,) dtypes.uchar

훈련과 속도

기가 커널 단위로 시간을 볼 수 있고, T4 colab에서 한 스텝이 약 75ms 정도예요. 왜 이렇게 느린지 언급되는데, tinygrad의 핵심 개념인 TinyJit 을 도입하면 커널을 캡처·재생해 훨씬 빨라집니다.

from tinygrad import TinyJit

jit_step = TinyJit(step)

데코레이터로도 쓸 수 있어요: @TinyJit. 함수의 처음 두 번 실행은 정상적으로 JIT가 커널을 캡처하고, 세 번째부터는 tinygrad 연산만 재생돼 파이썬 코드 실행 오버헤드를 건너뜁니다. 커널에 영향을 주는 non-tinygrad 파이썬 값은 두 번째 실행부터 "고정(frozen)"된다는 점을 알아두세요. Tensor 랜덤 함수는 예상대로 동작합니다.

큐 후 속도 경쟁은 with Context(BEAM=2) 로 BEAM 검색을 돌리면 가능한 구현들을 검색해 하드웨어에서 가장 빠른 걸 찾아주고, 결과가 캐시되니 다음에는 검색할 필요가 없어요. 보통 이 뒤엔 PyTorch와 속도 경쟁이 가능해집니다.

시작하기

tinygrad는 거의 순수 파이썬이고 짧아서 PyTorch와 달리 라이브러리 버그를 고치는 것도 충분히 가능해요. 이 예시에 멀티 GPU 지원을 추가하는 건 두 줄이면 됩니다 — 모델을 모든 GPU에 .shard 하고 데이터셋을 배치로 .shard 하면 돼요.

더 알아보기