Trax 빠른 시작 — fastmath·레이어·데이터

Trax 빠른 시작

Trax에 흐르는 기본 단위는 텐서(다차원 배열)예요. Trax는 numpy API를 쓰고, GPU·TPU로 가속하기 위해 trax.fastmath 패키지가 JAX와 TensorFlow NumPy 백엔드로 numpy 연산을 빠르게 수행하며 그래디언트를 자동 계산해요.

출처: https://trax-ml.readthedocs.io/en/latest/notebooks/trax_intro.html

fastmath로 텐서 계산

from trax.fastmath import numpy as fastnp

trax.fastmath.use_backend('jax')  # 'jax' 또는 'tensorflow-numpy'

matrix = fastnp.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
vector = fastnp.ones(3)
product = fastnp.dot(vector, matrix)
print(f'product = {product}')
tanh = fastnp.tanh(product)
print(f'tanh(product) = {tanh}')

그래디언트는 trax.fastmath.grad로 계산해요.

def f(x):
    return 2.0 * x * x

grad_f = trax.fastmath.grad(f)
print(f'grad(2x^2) at 1 = {grad_f(1.0)}')
print(f'grad(2x^2) at -2 = {grad_f(-2.0)}')

grad(2x^2) at 1 = 4.0
grad(2x^2) at -2 = -8.0

레이어 (Layers)

레이어는 Trax 모델의 기본 구성 블록이에요. 예를 들어 Embedding 레이어는 jnp.take로 토큰 ID를 벡터로 매핑하죠.

from trax import layers as tl

x = np.arange(15)
embedding = tl.Embedding(vocab_size=20, d_feature=32)
embedding.init(trax.shapes.signature(x))
y = embedding(x)
print(f'shape of y = {y.shape}')

shape of y = (15, 32)

모델 만들기 (Serial·Branch)

모델은 Serial과 Branch 콤비네이터로 레이어를 쌓아 만들어요. 감성 분류 모델 예시:

model = tl.Serial(
    tl.Embedding(vocab_size=8192, d_feature=256),
    tl.Mean(axis=1),       # axis 1 (문장 길이) 평균
    tl.Dense(2),           # 2 클래스 분류
)
print(model)

Serial[
  Embedding_8192_256
  Mean
  Dense_2
]

데이터

Trax에서 데이터 스트림은 파이썬 반복자로 표현돼요. next(data_stream)으로 (inputs, targets) 튜플을 받을 수 있고, TensorFlow Datasets를 쉽게 쓰거나 텍스트 파일에서 반복자를 얻을 수 있어요.

train_stream = trax.data.TFDS('imdb_reviews', keys=('text', 'label'), train=True)()
eval_stream = trax.data.TFDS('imdb_reviews', keys=('text', 'label'), train=False)()
print(next(train_stream))