Train neural networks

Who this path is for

Learners who can trace shapes, differentiate a scalar loss and carry state through an update.

Preparation

Shared foundations

Your first useful artifact

A CPU classifier with stable loss, replay checks and count-weighted evaluation.

How to study

Build the classifier first, freeze its evaluation split and add optimizer, random and data-position state before rehearsing interruption. Continue through CPU performance and deployment to the image harness. Use the reproducible-experiments workshop to compare changes. The optional language branch starts at Transformers after recovery; add pretraining and supervised adaptation next. Complete the RL phase before learned-reward PPO. These branches are separate from the classifier evidence.

Course phases and lessons

1. Setup & first steps

   Run a saved program and explain its environment, array output and device placement. Start on one CPU, practice sharding on four logical CPU devices, and learn how to get a TPU running, launch a TPU job, and run course experiments on TPU in parallel from Phase 00 onward.

   • Set up your learning workspace
   • Meet your arrays and devices
   • Practice with four virtual CPU devices
   • Move your experiment to a TPU
2. Arrays & pure functions

   Treat every axis as part of a contract. Work through normalization, broadcasting and immutable updates, then diagnose mistakes that still return plausible numbers.

   • From NumPy to jax.numpy
   • Shapes, broadcasting, and dtypes
   • Pure functions and explicit inputs
   • Immutable updates and indexing
3. JAX transformations

   Begin with a scalar function, then differentiate it, batch it and compile it. Compare each transformed result with a simple reference so you can explain what changed.

   • Your first gradient
   • Losses and value_and_grad
   • Batch a function with vmap
   • Compile a function with jit
   • Tracing, static arguments, and recompilation
4. State, randomness & control flow

   Make randomness, parameters and loop state explicit. Learn how to repeat an experiment without accidentally reusing randomness or mutating a caller’s state.

   • Random keys without surprises
   • Pytrees and structured parameters
   • Compiled loops with lax.scan
   • Branches with lax.cond
5. Math & optimization

   Connect geometry and probability to the updates in a training loop. Begin with linear algebra, gradient checks and gradient descent, then investigate conditioning, curvature, likelihoods, regularization and optimizer behavior.

   • Vectors, norms, and projections
   • Linear algebra and loss intuition
   • Least squares, rank, and conditioning
   • The chain rule, Jacobians, and directions
   • Gradient checking and numerical accuracy
   • Curvature, Hessians, and learning rates
   • Write gradient descent yourself
   • Probability, likelihood, and stable losses
   • Regularization and honest validation
   • Minibatches, expectation, and gradient noise
   • Optimize with Optax
   • Adam, clipping, and learning-rate schedules
6. Neural network training

   Build the same ideas into small neural networks. Understand hidden layers, parameter state, stable objectives and evaluation before moving to a small image classifier.

   • Build a tiny multilayer perceptron
   • Model and state with Flax NNX
   • A compiled train and evaluation step
   • Train a CNN on a small dataset
   • Debug unstable learning
7. Data & checkpoint recovery

   Follow examples from the dataset into a training update, then save enough information to reproduce the next update after interruption. Track experiment configuration, data identity and metric histories with MLflow alongside complete checkpoints.

   • Design an input pipeline
   • Load, batch, and checkpoint order with Grain
   • Save model, optimizer, and random state
   • Recover data position and resume
   • Asynchronous checkpoints and failure boundaries
   • Track experiments and lineage with MLflow
8. Performance diagnosis

   Measure completed work, inspect a real profiler trace and relate compiler output to memory traffic. Compare rematerialization choices while separating saved residuals, compiler memory estimates and observed latency.

   • Benchmark asynchronous work correctly
   • Diagnose recompilation and host synchronization
   • Read an XProf trace
   • Reason with HLO and the roofline model
   • Memory, rematerialization, and optimization tradeoffs
9. Deployment, interoperability & edge AI

   Adapt a trained model, verify framework and export contracts, compare precision policies and measure local inference before planning server or edge capacity. Package a model service in a reproducible container and verify readiness, artifact identity and request contracts.

   • Adapt a pretrained model and choose a post-training objective
   • Move models between JAX, Keras, TensorFlow and PyTorch
   • Convert PyTorch weights to Flax and locate numerical errors
   • Export a computation and verify its serving contract
   • Choose weight, activation and accumulation precision
   • Containerize a model service and verify its boundary
   • Inference capacity, batching, and autoscaling
   • Deploy at the edge: conversion, budgets and device checks

A recoverable image classifier with verified exported inference

A guided project is available for this route.

Complete the eight image lifecycle stages and changed-condition synthesis: independent CNN checks, full-state recovery, frozen error analysis, precision, export and measured inference. Transformer-specific evidence belongs to the complementary text harness.

• Training and evaluation splits, seeded configuration and learning curves
• A saved full state whose next update matches an uninterrupted run
• Predictions from the restored model and examples explaining its failures

Read the portfolio project guide

Work through the synthesis assessment

Use this path in your learning workspace