JAX là framework cho phép viết mã Python có khả năng biên dịch sang XLA, tối ưu cho cả CPU và TPU. Dưới đây là cách xây dựng một mạng nơ-ron đa tầng để phân loại ảnh chữ số viết tay.

Cấu trúc mạng nơ-ron

Mạng gồm ba phần chính: hàm khởi tạo trọng số, hàm truyền tiếp (forward pass), và cơ chế tối ưu gradient. JAX sử dụng cơ chế pure function, nghĩa là mọi tham số đều phải được truyền rõ ràng.

import jax.numpy as jnp
from jax import grad, jit, vmap
from jax import random as jrand

# Hàm kích hoạt và chuẩn hóa đầu ra
def activate(x):
    return jnp.maximum(0.0, x)  # ReLU

def normalize_logits(x):
    exp_x = jnp.exp(x - jnp.max(x, axis=-1, keepdims=True))
    return exp_x / jnp.sum(exp_x, axis=-1, keepdims=True)

# Truyền tiếp qua một tầng
def layer_forward(inp, weights, bias):
    return jnp.dot(inp, weights) + bias

# Mô hình đầy đủ
def classifier(model_weights, inp_data):
    activations = inp_data
    num_layers = len(model_weights)
    
    for idx, (wt, bs) in enumerate(model_weights):
        linear = layer_forward(activations, wt, bs)
        # Tầng cuối dùng softmax, các tầng khác dùng ReLU
        if idx < num_layers - 1:
            activations = activate(linear)
        else:
            activations = normalize_logits(linear)
    
    return activations

Khởi tạo tham số mô hình

Trọng số được khởi tạo theo phân phối chuẩn với độ lệch chuẩn nhỏ để tránh vấn đề gradient bão hòa hoặc tiêu biến.

def create_weights(rng_key, dims_list):
    keys = jrand.split(rng_key, len(dims_list))
    params = []
    for k, (in_d, out_d) in zip(keys, zip(dims_list[:-1], dims_list[1:])):
        scale = jnp.sqrt(2.0 / in_d)  # Khởi tạo He
        w = jrand.normal(k, (in_d, out_d)) * scale
        b = jnp.zeros(out_d)
        params.append((w, b))
    return params

# Định nghĩa kiến trúc
feature_count = 784   # Ảnh 28x28
hidden_units = 256
class_count = 10

seed = jrand.PRNGKey(42)
model_dims = [feature_count, hidden_units, class_count]
network = create_weights(seed, model_dims)

Hàm mất mát và đánh giá

Sử dụng cross-entropy chuẩn hóa làm hàm mất mát. JAX cung cấp vmap để vector hóa phép tính trên toàn bộ batch mà không cần vòng lặp thủ công.

def categorical_loss(weights, batch_x, batch_y):
    logits = classifier(weights, batch_x)
    # Tránh log(0) bằng cách thêm epsilon
    stable_logits = jnp.clip(logits, 1e-7, 1.0)
    log_probs = jnp.log(stable_logits)
    return -jnp.mean(jnp.sum(batch_y * log_probs, axis=-1))

def compute_accuracy(weights, data_x, data_y):
    predictions = classifier(weights, data_x)
    pred_classes = jnp.argmax(predictions, axis=-1)
    true_classes = jnp.argmax(data_y, axis=-1)
    return jnp.mean(pred_classes == true_classes)

Vòng lặp huấn luyện với Adam tùy chỉnh

Thay vì dùng jax.experimental.optimizers (đã lỗi thời), ta tự triển khai Adam để kiểm soát tốt hơn:

def init_adam(params):
    m_t = [[jnp.zeros_like(w), jnp.zeros_like(b)] for w, b in params]
    v_t = [[jnp.zeros_like(w), jnp.zeros_like(b)] for w, b in params]
    return m_t, v_t

@jit
def adam_update(weights, grads, m_state, v_state, step, lr=1e-3, b1=0.9, b2=0.999, eps=1e-8):
    m_new, v_new, w_new = [], [], []
    
    for (w, b), (dw, db), (m_w, m_b), (v_w, v_b) in zip(weights, grads, m_state, v_state):
        # Cập nhật moment bậc nhất
        m_w_upd = b1 * m_w + (1 - b1) * dw
        m_b_upd = b1 * m_b + (1 - b1) * db
        
        # Cập nhật moment bậc hai
        v_w_upd = b2 * v_w + (1 - b2) * (dw ** 2)
        v_b_upd = b2 * v_b + (1 - b2) * (db ** 2)
        
        # Hiệu chỉnh bias
        m_w_hat = m_w_upd / (1 - b1 ** step)
        m_b_hat = m_b_upd / (1 - b1 ** step)
        v_w_hat = v_w_upd / (1 - b2 ** step)
        v_b_hat = v_b_upd / (1 - b2 ** step)
        
        # Cập nhật trọng số
        w_new.append((w - lr * m_w_hat / (jnp.sqrt(v_w_hat) + eps),
                      b - lr * m_b_hat / (jnp.sqrt(v_b_hat) + eps)))
        
        m_new.append((m_w_upd, m_b_upd))
        v_new.append((v_w_upd, v_b_upd))
    
    return w_new, m_new, v_new

Luồng huấn luyện chính

def training_loop(key, weights, train_batches, epochs=15, batch_sz=64):
    m, v = init_adam(weights)
    
    for ep in range(epochs):
        key, subkey = jrand.split(key)
        # Xáo trộn dữ liệu (giả định đã có hàm shuffle)
        
        epoch_loss = 0.0
        for batch_x, batch_y in train_batches:
            # Tính gradient
            loss_val, grads = jax.value_and_grad(categorical_loss)(weights, batch_x, batch_y)
            
            # Cập nhật Adam
            weights, m, v = adam_update(weights, grads, m, v, ep + 1)
            epoch_loss += loss_val
        
        print(f"Epoch {ep+1:02d} | Loss trung bình: {epoch_loss / len(train_batches):.4f}")
    
    return weights

Lưu trữ và triển khai mô hình

Sau khi huấn luyện, lưu mô hình bằng jax.tree_util và pickle:

import pickle

def serialize_model(weights, filepath):
    with open(filepath, 'wb') as f:
        pickle.dump(weights, f)

def load_model(filepath):
    with open(filepath, 'rb') as f:
        return pickle.load(f)

API phục vụ dự đoán với FastAPI

Thay vì Flask, dùng FastAPI để xử lý bất đồng bộ tốt hơn:

from fastapi import FastAPI, UploadFile, File
from PIL import Image
import io

api = FastAPI(title="Nhận diện chữ số JAX")

# Load mô hình khi khởi động
inference_model = load_model('digit_model.pkl')

@api.post("/infer")
async def predict_digit(file: UploadFile = File(...)):
    # Đọc và tiền xử lý ảnh
    contents = await file.read()
    img = Image.open(io.BytesIO(contents)).convert('L').resize((28, 28))
    img_arr = jnp.array(img, dtype=jnp.float32) / 255.0
    img_flat = img_arr.reshape(1, -1)  # Định dạng batch đơn lẻ
    
    # Dự đoán
    probs = classifier(inference_model, img_flat)
    predicted = int(jnp.argmax(probs, axis=-1)[0])
    confidence = float(jnp.max(probs))
    
    return {"digit": predicted, "confidence": confidence}

Chạy server với Uvicorn: uvicorn main:api --reload --port 8080