Deep Learning with JAX
Grigory Sapunov
- Editorial
- Manning
- Año
- 2024
- Páginas
- ~410
- Nivel
- Intermedio
De qué trata
Manual exhaustivo y moderno que introduce al desarrollador en el paradigma de programación funcional y la aceleración de hardware para inteligencia artificial con JAX, la biblioteca de cálculo numérico de Google. La obra está dirigida a ingenieros de machine learning, investigadores de deep learning y desarrolladores de Python de nivel intermedio a avanzado que ya cuentan con experiencia en bibliotecas como PyTorch, TensorFlow o NumPy, y buscan dominar una tecnología que escala eficientemente en clústeres de GPUs y TPUs.
El autor enseña la filosofía de pureza funcional e inmutabilidad de JAX, la gestión determinista de números pseudoaleatorios con claves explícitas (`jax.random`), la estructuración de contenedores con pytrees (`jax.tree_util`) y las cuatro transformaciones fundamentales que sustentan el framework: diferenciación automática inversa y hacia adelante (`jax.grad` y `jax.value_and_grad`), compilación Just-in-Time mediante el compilador XLA (`jax.jit`), vectorización por lotes (`jax.vmap`) y computación paralela distribuida (`jax.pmap`, `NamedSharding` y mallas de cómputo). Adicionalmente, examina el ecosistema de librerías asociadas: desarrollo declarativo de redes con Flax (módulos Linen y contenedor TrainState), optimizadores con Optax, utilidades de Haiku y Equinox, y el guardado de checkpoints distribuidos con Orbax.
A lo largo del libro se implementan modelos progresivos mediante cuadernos ejecutables en Google Colab y Cloud TPUs: comenzando por un clasificador de imágenes MNIST escrito desde cero en JAX puro (manejando manualmente los gradientes y el estado funcional de los parámetros), continuando con redes convolucionales tipo ResNet utilizando Flax y Optax, y finalizando con la integración de arquitecturas Transformer autorregresivas (estilo GPT) y modelos de difusión para generación de imágenes (Stable Diffusion) a través del ecosistema de Hugging Face. El libro aporta una claridad excepcional para desaprender los hábitos orientados a objetos de PyTorch y adoptar el pensamiento funcional optimizado para compilación; conviene tener bases matemáticas de cálculo multivariable y no es apto para principiantes absolutos en programación.
- JAX
- Flax
- Optax
- Compilación JIT
- XLA
- Diferenciación automática
- Vectorización vmap
- Tensor Sharding
- TPUs y GPUs
- Programación funcional
- Transformers
¿Cómo enseña?
Con proyectos
- ✗ Sin ejercicios
- ✓ Incluye proyectos
Proyectos
Clasificador MNIST desde cero en JAX puro
Red neuronal feedforward programada sin dependencias de alto nivel, implementando bucles de entrenamiento funcionales, actualización de parámetros en pytrees y transformaciones con grad, jit y vmap.
Red convolucional profunda ResNet con Flax y Optax
Arquitectura moderna de visión computacional para clasificación de imágenes construida con módulos de Flax Linen, seguimiento de estado con TrainState y optimización de gradientes mediante Optax.
Pipeline generativo con Transformers y modelos de difusión
Implementación práctica de inferencia y afinamiento de modelos de lenguaje autorregresivos tipo GPT y síntesis de imágenes con Stable Diffusion integrando librerías de Hugging Face en JAX.
