آموزش JAX: برنامه‌نویسی تابعی و شتاب‌دهنده‌ها در یادگیری عمیق

بسیاری از توسعه‌دهندگان یادگیری ماشین کار خود را با فریم‌ورک‌های متداول مثل PyTorch آغاز می‌کنند؛ جایی که همه‌چیز در قالب شیءگرایی، کلاس‌های `nn.Module` و تغییر مستقیم وضعیت تنسورها تعریف می‌شود. با این حال، هنگام مقیاس‌پذیری به خوشه‌های عظیم سخت‌افزاری، سربار مفسر پایتون و ردیابی تک‌به‌تک عملیات می‌تواند کارایی محاسبات را به شدت کاهش دهد.

فریم‌ورک JAX این مشکل را با بازتعریف یادگیری عمیق در قالب برنامه‌نویسی تابعی حل کرده است. در این ساختار، خبری از کلاس‌ها و تغییر وضعیت پنهان نیست؛ بلکه مدل‌ها مجموعه‌ای از توابع خالص هستند که به کامپایلر تحویل داده می‌شوند تا مستقیماً به کد ماشین بهینه‌سازی‌شده برای GPU و TPU تبدیل شوند.

اگر می‌خواهید درک کنید چگونه غول‌های هوش مصنوعی مانند DeepMind و Anthropic مدل‌هایی مثل Gemini و Claude را آموزش می‌دهند، یادگیری JAX و شیوه‌ی تفکر تابعی آن گامی ضروری است. در این درس با مفاهیم پایه‌ای JAX نظیر `grad`، `jit`، `vmap` و ساختار پای‌تری آشنا می‌شوید و تفاوت بنیادین آن را با رویکرد شیءگرای PyTorch خواهید دید.

فریم‌ورک PyTorch تنسورها را تغییر می‌دهد (mutate). فریم‌ورک TensorFlow گراف می‌سازد. فریم‌ورک JAX توابع خالص را کامپایل می‌کند. همین ویژگیِ آخر، طرز فکر شما را درباره‌ی یادگیری عمیق دگرگون می‌کند.

اهداف یادگیری

• نوشتن کد شبکه‌ی عصبی بر پایه‌ی توابع خالص با استفاده از رابط تابعی JAX (شامل jax.numpy، jax.grad، jax.jit و jax.vmap)
• توضیح تفاوت کلیدی در طراحی بین اجرای فوری و تغییرپذیر PyTorch (eager mutation) و مدل کامپایل تابعی JAX
• به‌کارگیری کامپایل jit و برداری‌سازی با vmap برای شتاب‌بخشی به حلقه‌های آموزش در مقایسه با پایتون خام
• آموزش یک شبکه‌ی ساده در JAX و مقایسه‌ی مدیریت صریح وضعیت (explicit state management) با رویکرد شیءگرای PyTorch

صورت مسئله

نحوه‌ی ساخت شبکه‌های عصبی در PyTorch را بلد هستید: یک nn.Module تعریف می‌کنید، متد .backward() را فراخوانی می‌کنید و با بهینه‌ساز (optimizer) یک گام برمی‌دارید. این روش کار می‌کند و میلیون‌ها نفر هم از آن استفاده می‌کنند.

اما PyTorch محدودیتی ساختاری دارد که در دی‌ان‌ای آن ریشه دوانده است: عملیات را به‌صورت فوری (eagerly)، یکی پس از دیگری در پایتون ردگیری (trace) می‌کند. هر tensor + tensor یک اجرای کرنل جداگانه است. در هر گامِ آموزش، همان کد پایتون دوباره تفسیر می‌شود. این سازوکار تا زمانی که نخواهید یک مدل 540 میلیارد پارامتری را روی 2048 پردازنده‌ی TPU آموزش دهید، خوب کار می‌کند؛ اما در آن مقیاس، این سربار شما را از پا درمی‌آورد.

شرکت Google DeepMind مدل Gemini را روی JAX آموزش می‌دهد. شرکت Anthropic نیز مدل Claude را با JAX آموزش داده است. این پروژه‌ها کارهای کوچکی نیستند؛ بلکه بزرگ‌ترین دوره‌های آموزش شبکه‌ی عصبی روی کره‌ی زمین به شمار می‌روند. آن‌ها JAX را انتخاب کردند چون با حلقه‌ی آموزش مانند یک برنامه‌ی قابل کامپایل رفتار می‌کند، نه دنباله‌ای از فراخوانی‌های پایتون.

فریم‌ورک JAX همان NumPy است، ولی با سه ابرقدرت: مشتق‌گیری خودکار، کامپایل درجا (JIT) به XLA و برداری‌سازی خودکار. شما تابعی می‌نویسید که یک نمونه‌ی تکی را پردازش می‌کند؛ سپس JAX تابعی تحویلتان می‌دهد که یک بچ (batch) را پردازش می‌کند، گرادیان‌ها را به دست می‌آورد، به کد ماشین کامپایل می‌شود و روی چندین دستگاه اجرا می‌گردد؛ همه‌ی این‌ها بدون کوچک‌ترین تغییری در تابع اصلی.

مفاهیم اصلی

فلسفه‌ی JAX

فریم‌ورک JAX مبتنی بر پارادایم تابعی است. در آن خبری از کلاس‌ها، وضعیت تغییرپذیر و متد .backward() نیست. در عوض:

این موضوع صرفاً یک سلیقه‌ی نگارشی نیست؛ یک الزام برای کامپایلر است. کامپایل JIT به توابع خالص نیاز دارد؛ یعنی توابعی که ورودی‌های یکسان در آن‌ها همواره خروجی‌های یکسان تولید کنند و فاقد هرگونه اثر جانبی (side effect) باشند. دقیقاً همین محدودیت است که افزایش سرعت تا 100 برابر را ممکن می‌سازد.

jax.numpy: ظاهر آشنا

فریم‌ورک JAX رابط برنامه‌نویسی NumPy را برای شتاب‌دهنده‌ها بازپیاده‌سازی کرده است:

import jax.numpy as jnp

a = jnp.array([1.0, 2.0, 3.0])
b = jnp.array([4.0, 5.0, 6.0])
c = jnp.dot(a, b)

همان نام توابع، همان قوانین برادکستینگ (broadcasting) و همان رفتار برش‌زنی (slicing) برقرار است؛ با این تفاوت که آرایه‌ها روی GPU یا TPU قرار دارند و تک‌تک عملیات برای کامپایلر قابل ردگیری است.

یک تفاوت حیاتی وجود دارد: آرایه‌های JAX تغییرناپذیرند. یعنی نمی‌توانید بنویسید a[0] = 5؛ بلکه باید بنویسید: a = a.at[0].set(5). این شیوه شاید در هفته‌ی اول غریب به نظر برسد، اما بعد جا می‌افتد: تغییرناپذیری همان عاملی است که ترکیب تحول‌هایی مانند grad، jit و vmap را ممکن می‌سازد.

jax.grad: مشتق‌گیری خودکار تابعی

فریم‌ورک PyTorch گرادیان‌ها را به تنسورها متصل می‌کند (.grad). فریم‌ورک JAX گرادیان‌ها را به توابع متصل می‌کند.

import jax

def f(x):
    return x ** 2

df = jax.grad(f)
df(3.0)

متد jax.grad تابعی را می‌گیرد و تابع جدیدی برمی‌گرداند که گرادیان را حساب می‌کند. نه فراخوانی .backward() در کار است و نه گراف محاسباتی روی تنسورها ذخیره می‌شود. گرادیان فقط یک تابع دیگر است که می‌توانید آن را فراخوانی کنید، با بقیه ترکیب کنید یا با JIT کامپایلش نمایید.

این ترکیب‌پذیری هیچ محدودیتی ندارد:

d2f = jax.grad(jax.grad(f))
d2f(3.0)

مشتق‌های مرتبه‌ی دوم، مشتق‌های مرتبه‌ی سوم، ژاکوبی‌ها و هسی‌ها؛ همگی تنها با ترکیب کردن grad به دست می‌آیند. فریم‌ورک PyTorch هم می‌تواند این کار را انجام دهد (torch.autograd.functional.hessian)، اما این قابلیت بعداً به آن چسبانده شده است؛ در حالی که در JAX، پایه‌ی اصلی فریم‌ورک محسوب می‌شود.

محدودیت اصلی: grad فقط روی توابع خالص کار می‌کند. داخل تابع نباید از دستورهای print استفاده کنید (چون موقع ردگیری اجرا می‌شوند، نه موقع اجرا). نباید وضعیت خارجی را دستکاری کنید و بدون مدیریت صریح کلیدها، حق تولید اعداد تصادفی را ندارید.

jit: کامپایل به XLA

@jax.jit
def train_step(params, x, y):
    loss = loss_fn(params, x, y)
    return loss

fast_step = jax.jit(train_step)

در اولین فراخوانی، JAX تابع را ردگیری می‌کند؛ یعنی ثبت می‌کند چه عملیاتی رخ می‌دهند، بدون آنکه واقعاً آن‌ها را اجرا کند. سپس این رد را به XLA (مخفف Accelerated Linear Algebra)، یعنی کامپایلر گوگل برای TPU و GPU می‌سپارد. کامپایلر XLA عملیات را در هم ادغام می‌کند (fusion)، کپی‌های اضافی حافظه را حذف می‌نماید و کد ماشین بهینه‌شده تولید می‌کند.

فراخوانی‌های بعدی اصلاً وارد پایتون نمی‌شوند؛ کد کامپایل‌شده روی شتاب‌دهنده با سرعت ++C اجرا می‌شود.

مواردی که JIT به کار می‌آید:

• گام‌های آموزش (تکرار هزاران‌باره‌ی یک محاسبه‌ی یکسان)
• استنتاج (همان مدل، ولی با ورودی‌های متفاوت)
• هر تابعی که بیش از یک بار با ورودی‌هایی با ابعاد مشابه فراخوانی شود

مواردی که JIT کارایی را خراب می‌کند:

• توابعی با جریان کنترل پایتونیِ وابسته به مقادیر (مثل if x > 0 وقتی x یک آرایه‌ی ردگیری‌شده است)
• محاسباتی که فقط یک بار اجرا می‌شوند (هزینه‌ی کامپایل از زمان اجرا بیشتر می‌شود)
• خطایابی (مرحله‌ی ردگیری، اجرای واقعی را پنهان می‌کند)

محدودیت‌های مربوط به جریان کنترل کاملاً جدی هستند. در این فریم‌ورک، jax.lax.cond جایگزین if/else می‌شود و jax.lax.scan جای حلقه‌های for را می‌گیرد. استفاده از این ساختارها اختیاری نیست؛ بهایی است که باید برای کامپایل بپردازید.

vmap: برداری‌سازی خودکار

تابعی می‌نویسید که یک نمونه‌ی تکی را پردازش می‌کند:

def predict(params, x):
    return jnp.dot(params['w'], x) + params['b']

سپس vmap آن را ارتقا می‌دهد تا یک بچ کامل را پردازش کند:

batch_predict = jax.vmap(predict, in_axes=(None, 0))

عبارت in_axes=(None, 0) یعنی: روی params بچ‌بندی انجام نده (مشترک است)، ولی روی بعد 0 ورودیِ x برداری‌سازی کن. دیگر به حلقه‌ی for دستی، تغییر ابعاد (reshaping) و سر و کله زدن با بعدِ بچ نیازی نیست. فریم‌ورک JAX بعد بچ را خودش تشخیص می‌دهد و کل محاسبات را برداری‌سازی می‌کند.

این صرفاً یک قند نحوی (syntactic sugar) نیست؛ vmap کد برداری‌شده‌ی ادغام‌شده‌ای تولید می‌کند که 10 تا 100 برابر سریع‌تر از حلقه‌ی پایتون اجرا می‌شود. علاوه بر این، با jit و grad نیز ترکیب‌پذیر است:

per_example_grads = jax.vmap(jax.grad(loss_fn), in_axes=(None, 0, 0))

گرادیان‌های به ازای هر نمونه (per-example gradients) تنها در یک خط. این کار در PyTorch تقریباً بدون حقه‌های عجیب‌وغریب ناممکن است.

pmap: موازی‌سازی داده در چندین دستگاه

parallel_step = jax.pmap(train_step, axis_name='devices')

متد pmap تابع را روی همه‌ی دستگاه‌های دردسترس (GPUها/TPUها) تکثیر می‌کند و بچ ورودی را بین آن‌ها تقسیم می‌نماید. در داخل تابع، متدهای jax.lax.pmean و jax.lax.psum گرادیان‌ها را بین دستگاه‌های مختلف همگام‌سازی می‌کنند.

گوگل مدل Gemini را روی هزاران تراشه‌ی TPU v5e با استفاده از pmap (و جانشین آن، shard_map) آموزش می‌دهد. مدل برنامه‌نویسی ساده است: نسخه‌ی تک‌دستگاهی را بنویسید، آن را داخل pmap بپیچید و کار تمام است.

پای‌تری‌ها: ساختار داده‌ی همگانی

فریم‌ورک JAX روی «پای‌تری‌ها» (pytrees) کار می‌کند؛ یعنی ساختارهای تودرتویی از لیست‌ها، تاپل‌ها، دیکشنری‌ها و آرایه‌ها. پارامترهای مدل شما یک پای‌تری هستند:

params = {
    'layer1': {'w': jnp.zeros((784, 256)), 'b': jnp.zeros(256)},
    'layer2': {'w': jnp.zeros((256, 128)), 'b': jnp.zeros(128)},
    'layer3': {'w': jnp.zeros((128, 10)),  'b': jnp.zeros(10)},
}

تمام تحول‌های JAX—مانند grad، jit و vmap—می‌دانند چطور روی ساختارهای پای‌تری حرکت کنند. تابع jax.tree.map(f, tree) تابع f را روی تک‌تک برگ‌ها اجرا می‌کند. بهینه‌سازها دقیقاً با همین روش تمام پارامترها را یک‌جا به‌روزرسانی می‌کنند:

params = jax.tree.map(lambda p, g: p - lr * g, params, grads)

دیگر نیازی به متد .parameters() و ثبت پارامترها نیست. خود ساختار درختی، همان مدل است.

مقایسه‌ی رویکرد تابعی با شیءگرا

فریم‌ورک PyTorch وضعیت را داخل شیء نگه می‌دارد:

class Model(nn.Module):
    def __init__(self):
        self.linear = nn.Linear(784, 10)

    def forward(self, x):
        return self.linear(x)

در مقابل، JAX از توابع خالص با وضعیت صریح استفاده می‌کند:

def predict(params, x):
    return jnp.dot(x, params['w']) + params['b']

پارامترها به عنوان ورودی پاس داده می‌شوند. هیچ چیزی ذخیره نمی‌شود و هیچ داده‌ای تغییر نمی‌کند. همین ویژگی باعث می‌شود تک‌تک توابع آزمون‌پذیر، ترکیب‌پذیر و قابل کامپایل باشند. این ساختار همچنین یعنی مدیریت پارامترها تماماً به عهده‌ی خود شماست—یا اینکه باید به سراغ کتابخانه‌هایی مانند Flax یا Equinox بروید.

زیست‌بوم JAX

فریم‌ورک JAX ابزارهای پایه را در اختیارتان می‌گذارد؛ کتابخانه‌ها ارگونومی و راحتی کار را فراهم می‌کنند:

کتابخانه‌ی Optax ابزار استاندارد برای بهینه‌سازهاست. این کتابخانه تحول گرادیان (Adam، SGD، برش گرادیان) را از فرایند به‌روزرسانی پارامتر جدا می‌کند و ترکیب آن‌ها را به کاری بسیار ساده تبدیل می‌نماید:

optimizer = optax.chain(
    optax.clip_by_global_norm(1.0),
    optax.adam(learning_rate=1e-3),
)

چه زمانی از JAX استفاده کنیم و چه زمانی از PyTorch؟

پاسخ صادقانه این است: تا زمانی که دلیل مشخصی برای استفاده از JAX ندارید، با همان PyTorch کار کنید. آن دلایل مشخص عبارتند از: دسترسی به TPU، نیاز به گرادیان‌های به ازای هر نمونه، آموزش روی چندین دستگاه در مقیاس بسیار بزرگ، یا کار در شرکت‌هایی مانند Google، DeepMind و Anthropic.

اعداد تصادفی در JAX

در JAX هیچ وضعیت تصادفی سراسری وجود ندارد. هر عملیات تصادفی به یک کلید صریح PRNG (تولیدکننده‌ی اعداد شبه‌تصادفی) نیاز دارد:

key = jax.random.PRNGKey(42)
key1, key2 = jax.random.split(key)
w = jax.random.normal(key1, shape=(784, 256))

این سازوکار شاید اوایل آزاردهنده به نظر برسد؛ اما بازتولیدپذیری (reproducibility) را در تمامی دستگاه‌ها و دفعات کامپایل تضمین می‌کند؛ قابلیتی که متد torch.manual_seed در PyTorch نمی‌تواند در محیط‌های چند GPU تضمین کند.

پیاده‌سازی عملی

گام 1: راه‌اندازی و داده‌ها

یک پرسپترون چندلایه (MLP) سه‌لایه را با استفاده از JAX و Optax روی MNIST آموزش می‌دهیم: 784 ورودی، دو لایه‌ی پنهان با 256 و 128 نورون، و 10 کلاس خروجی.

import jax
import jax.numpy as jnp
from jax import random
import optax

def get_mnist_data():
    from sklearn.datasets import fetch_openml
    mnist = fetch_openml('mnist_784', version=1, as_frame=False, parser='auto')
    X = mnist.data.astype('float32') / 255.0
    y = mnist.target.astype('int')
    X_train, X_test = X[:60000], X[60000:]
    y_train, y_test = y[:60000], y[60000:]
    return X_train, y_train, X_test, y_test

گام 2: مقداردهی اولیه‌ی پارامترها

خبری از کلاس نیست؛ تنها تابعی می‌نویسیم که یک پای‌تری برگرداند:

def init_params(key):
    k1, k2, k3 = random.split(key, 3)
    scale1 = jnp.sqrt(2.0 / 784)
    scale2 = jnp.sqrt(2.0 / 256)
    scale3 = jnp.sqrt(2.0 / 128)
    params = {
        'layer1': {
            'w': scale1 * random.normal(k1, (784, 256)),
            'b': jnp.zeros(256),
        },
        'layer2': {
            'w': scale2 * random.normal(k2, (256, 128)),
            'b': jnp.zeros(128),
        },
        'layer3': {
            'w': scale3 * random.normal(k3, (128, 10)),
            'b': jnp.zeros(10),
        },
    }
    return params

مقداردهی اولیه به روش He، به‌صورت دستی پیاده شده است. سه کلید PRNG از یک سید (seed) اولیه مشتق شده‌اند. تک‌تک وزن‌ها آرایه‌هایی تغییرناپذیر در یک دیکشنری تودرتو هستند.

گام 3: گذر رو به جلو

def forward(params, x):
    x = jnp.dot(x, params['layer1']['w']) + params['layer1']['b']
    x = jax.nn.relu(x)
    x = jnp.dot(x, params['layer2']['w']) + params['layer2']['b']
    x = jax.nn.relu(x)
    x = jnp.dot(x, params['layer3']['w']) + params['layer3']['b']
    return x

def loss_fn(params, x, y):
    logits = forward(params, x)
    one_hot = jax.nn.one_hot(y, 10)
    return -jnp.mean(jnp.sum(jax.nn.log_softmax(logits) * one_hot, axis=-1))

توابع خالص: پارامترها وارد می‌شوند و پیش‌بینی خارج می‌شود؛ بدون self و بدون وضعیت ذخیره‌شده. تابع زیان (loss) با نام loss_fn آنتروپی متقاطع (Cross-Entropy) را از پایه حساب می‌کند: سافت‌مکس (Softmax)، لگاریتم و منفی میانگین.

گام 4: گام آموزش کامپایل‌شده با JIT

@jax.jit
def train_step(params, opt_state, x, y):
    loss, grads = jax.value_and_grad(loss_fn)(params, x, y)
    updates, opt_state = optimizer.update(grads, opt_state, params)
    params = optax.apply_updates(params, updates)
    return params, opt_state, loss

@jax.jit
def accuracy(params, x, y):
    logits = forward(params, x)
    preds = jnp.argmax(logits, axis=-1)
    return jnp.mean(preds == y)

تابع jax.value_and_grad هم‌زمان هم مقدار تابع زیان و هم گرادیان‌ها را در یک دور برمی‌گرداند. دکوراتور @jax.jit هر دو تابع را با XLA کامپایل می‌کند. پس از اولین فراخوانی، هر گام آموزش بدون دست زدن به پایتون اجرا می‌شود.

گام 5: حلقه‌ی آموزش

optimizer = optax.adam(learning_rate=1e-3)

X_train, y_train, X_test, y_test = get_mnist_data()
X_train, X_test = jnp.array(X_train), jnp.array(X_test)
y_train, y_test = jnp.array(y_train), jnp.array(y_test)

key = random.PRNGKey(0)
params = init_params(key)
opt_state = optimizer.init(params)

batch_size = 128
n_epochs = 10

for epoch in range(n_epochs):
    key, subkey = random.split(key)
    perm = random.permutation(subkey, len(X_train))
    X_shuffled = X_train[perm]
    y_shuffled = y_train[perm]

    epoch_loss = 0.0
    n_batches = len(X_train) // batch_size
    for i in range(n_batches):
        start = i * batch_size
        xb = X_shuffled[start:start + batch_size]
        yb = y_shuffled[start:start + batch_size]
        params, opt_state, loss = train_step(params, opt_state, xb, yb)
        epoch_loss += loss

    train_acc = accuracy(params, X_train[:5000], y_train[:5000])
    test_acc = accuracy(params, X_test, y_test)
    print(f"Epoch {epoch + 1:2d} | Loss: {epoch_loss / n_batches:.4f} | "
          f"Train Acc: {train_acc:.4f} | Test Acc: {test_acc:.4f}")

10 دوره (epoch) آموزش؛ حدود 97% دقت در داده‌های آزمون. دوره‌ی اول به خاطر کامپایل JIT کند است؛ دوره‌های 2 تا 10 بسیار سریع اجرا می‌شوند.

به چیزهایی که در این کد نیستند دقت کنید: خبری از .zero_grad()، .backward() و .step() نیست. کل به‌روزرسانی در قالب فراخوانیِ یک تابعِ ترکیب‌شده انجام می‌شود. گرادیان‌ها محاسبه می‌شوند، با Adam تغییر شکل می‌یابند و روی پارامترها اعمال می‌گردند؛ همگی داخل train_step.

کاربرد عملی

فریم‌ورک Flax: استاندارد گوگل

کتابخانه‌ی Flax رایج‌ترین کتابخانه‌ی شبکه‌ی عصبی در JAX است. این ابزار nn.Module را دوباره به کار می‌گیرد، اما با مدیریت صریح وضعیت:

import flax.linen as nn

class MLP(nn.Module):
    @nn.compact
    def __call__(self, x):
        x = nn.Dense(256)(x)
        x = nn.relu(x)
        x = nn.Dense(128)(x)
        x = nn.relu(x)
        x = nn.Dense(10)(x)
        return x

model = MLP()
params = model.init(jax.random.PRNGKey(0), jnp.ones((1, 784)))
logits = model.apply(params, x_batch)

همان ساختار PyTorch است، ولی params از خود مدل جداست. متد model.init() پارامترها را می‌سازد و model.apply(params, x) گذر رو به جلو (forward pass) را اجرا می‌کند. شیء مدل هیچ وضعیتی ندارد.

فریم‌ورک Equinox: جایگزین پایتونی

کتابخانه‌ی Equinox (اثر Patrick Kidger) مدل‌ها را به شکل پای‌تری بازنمایی می‌کند:

import equinox as eqx

model = eqx.nn.MLP(
    in_size=784, out_size=10, width_size=256, depth=2,
    activation=jax.nn.relu, key=jax.random.PRNGKey(0)
)
logits = model(x)

خود مدل یک پای‌تری است؛ بنابراین به .apply() نیازی نیست. پارامترها صرفاً برگ‌های مدل هستند. این ساختار به طرز فکر JAX نزدیک‌تر است.

کتابخانه‌ی Optax: بهینه‌سازهای ترکیب‌پذیر

کتابخانه‌ی Optax تحول گرادیان را از به‌روزرسانی جدا می‌کند:

schedule = optax.warmup_cosine_decay_schedule(
    init_value=0.0, peak_value=1e-3,
    warmup_steps=1000, decay_steps=50000
)

optimizer = optax.chain(
    optax.clip_by_global_norm(1.0),
    optax.adamw(learning_rate=schedule, weight_decay=0.01),
)

برش گرادیان (Gradient clipping)، گرم‌کردن نرخ یادگیری (learning rate warmup) و کاهش وزن (Weight decay)؛ همگی به‌صورت زنجیره‌ای از تحول‌ها در کنار هم چیده شده‌اند. هر تحول گرادیان‌ها را می‌بیند، آن‌ها را دستکاری می‌کند و به تحول بعدی می‌سپارد؛ بدون نیاز به یک کلاس بهینه‌ساز غول‌پیکر و یکپارچه.

تحویل پروژه

نصب:

pip install jax jaxlib optax flax

برای پشتیبانی از GPU:

pip install jax[cuda12]

برای TPU (در Google Cloud):

pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

نکات مهم برای کارایی:

• اولین فراخوانی JIT به دلیل کامپایل کند است. قبل از بنچمارک گرفتن، حتماً مرحله‌ی گرم‌کردن (warm up) را اجرا کنید.
• از حلقه‌های پایتون روی آرایه‌های JAX داخل JIT بپرهیزید. از jax.lax.scan یا jax.lax.fori_loop استفاده کنید.
• دستور jax.debug.print() داخل JIT کار می‌کند؛ دستور معمولی print() کار نمی‌کند.
• برای پیدا کردن گلوگاه‌ها از jax.profiler یا TensorBoard استفاده کنید؛ کامپایل XLA می‌تواند گلوگاه‌ها را پنهان کند.
• فریم‌ورک JAX به‌صورت پیش‌فرض 75% از حافظه‌ی GPU را از قبل تخصیص می‌دهد. برای غیرفعال کردن این ویژگی، مقدار XLA_PYTHON_CLIENT_PREALLOCATE=false را تنظیم کنید.

چک‌پوینت‌گیری:

import orbax.checkpoint as ocp
checkpointer = ocp.PyTreeCheckpointer()
checkpointer.save('/tmp/model', params)
restored = checkpointer.restore('/tmp/model')

خروجی‌های این درس:

• outputs/prompt-jax-optimizer.md -- پرامپتی برای انتخاب پیکربندی مناسب بهینه‌ساز در JAX
• outputs/skill-jax-patterns.md -- مهارتی شامل الگوهای تابعی در JAX

تمرین‌ها

1. لایه‌ی دراپ‌اوت (dropout) را به شبکه‌ی MLP اضافه کنید. در JAX دراپ‌اوت به کلید PRNG نیاز دارد؛ کلید را در طول گذر رو به جلو هدایت کنید و برای هر لایه‌ی دراپ‌اوت آن را بشکنید (split کنید). سپس دقت آزمون را با و بدون دراپ‌اوت مقایسه کنید.

2. با استفاده از jax.vmap گرادیان‌های به ازای هر نمونه را برای یک بچ شامل 32 تصویر MNIST محاسبه کنید. نُرم گرادیان (gradient norm) را برای هر نمونه به دست آورید. کدام نمونه‌ها بزرگ‌ترین گرادیان‌ها را دارند و چرا؟

3. تابع گذر رو به جلوی دستی را با یک تابع عمومی به نام mlp_forward(params, x) جایگزین کنید که برای هر تعداد لایه‌ای کار کند. از jax.tree.leaves برای تشخیص خودکار عمق شبکه استفاده کنید.

4. گام آموزش را با دکوراتور @jax.jit و بدون آن بنچمارک بگیرید. زمان اجرای 100 گام را برای هر حالت ثبت کنید. این افزایش سرعت روی سخت‌افزار شما چقدر است؟ سربار زمانی کامپایل در اولین فراخوانی چقدر است؟

5. برش گرادیان را با ترکیب optax.chain(optax.clip_by_global_norm(1.0), optax.adam(1e-3)) پیاده‌سازی کنید. مدل را با و بدون برش آموزش دهید. نمودار نُرم گرادیان را در طول آموزش رسم کنید تا اثر آن را ببینید.

اصطلاحات کلیدی

مطالعه‌ی بیشتر

• مستندات JAX به نشانی https://jax.readthedocs.io/ -- مستندات رسمی، با آموزش‌های فوق‌العاده برای grad، jit و vmap
• مقاله‌ی «JAX: composable transformations of Python+NumPy programs» (تألیف Bradbury و همکاران، 2018) -- مقاله‌ی اصلی که فلسفه‌ی طراحی را تشریح می‌کند
• مستندات Flax به نشانی https://flax.readthedocs.io/ -- کتابخانه‌ی شبکه‌ی عصبی گوگل برای JAX
• مقاله‌ی Patrick Kidger با عنوان «Equinox: neural networks in JAX via callable PyTrees and filtered transformations» (سال 2021) -- جایگزینی با ساختار پایتونی برای Flax
• کتابخانه‌ی DeepMind با عنوان «Optax: composable gradient transformation and optimisation» -- کتابخانه‌ی استاندارد بهینه‌سازها
• مقاله‌ی «You Don't Know JAX» (اثر Colin Raffel، سال 2020) -- راهنمایی کاربردی برای شناخت نکات و الگوهای JAX، از زبان یکی از نویسندگان T5