Project 50: Neural Network Framework

A self-contained deep-dive from the canonical Math from Foundations to Machine Learning curriculum.

  • Difficulty: Level 5: Master
  • Time: 4-8 weeks
  • Language: Python with NumPy
  • Prerequisites: Projects 44-49
  • Merged from: NN BrainInABox; ML-Found P20

What You Will Build

Turn the ad hoc components from earlier neural projects into a cohesive Mini-PyTorch/BrainInABox library. Define tensors with autograd, trainable parameters, nested modules, linear and convolutional layers, BatchNorm, activations, losses, Sequential, optimizers such as SGD and Adam, data loaders, train/eval modes, registered parameters and buffers, state dictionaries, serialization, and a reusable training loop. The public API should feel small and consistent while preserving correct parameter discovery, gradient accumulation, broadcasting, normalization state, and numerical stability. Prove the framework by rebuilding the earlier MLP and a CNN without model-specific training code.

Real World Outcome

A separate example program imports your package, prints a model summary and parameter count, trains an MLP on MNIST and a BatchNorm-equipped CNN on CIFAR-10, saves a checkpoint, loads it in a fresh process, and reproduces predictions. Training and evaluation runs demonstrate that BatchNorm uses batch statistics only while training and saved running statistics during evaluation. The test suite compares selected operations, gradients, buffers, and optimizer updates with trusted numerical/reference results.

Core Question

How do deep-learning frameworks organize differentiation, parameters, layers, optimization, data, and persistence into one safe composable system?

Concepts You Must Understand First

  • Dynamic autograd graphs and tensor operations — Project 44.
  • Module composition and parameter ownership — PyTorch nn.Module design as a reference.
  • Dense/convolutional forward and backward contracts — Projects 45 and 48.
  • Batch normalization: per-channel statistics, learnable scale/shift, running mean/variance buffers, epsilon, and distinct train/eval behavior.
  • Optimizer state, momentum, Adam moments, and zero-gradient lifecycle.
  • Python data model/operator overloading — Fluent Python by Luciano Ramalho.

Build Milestones

  1. Stabilize tensor/autograd behavior, broadcasting rules, dtypes, shape validation, and gradient tests.
  2. Implement Parameter, Module, recursive parameter registration, Linear, activations, losses, and containers.
  3. Add SGD/Adam, data loader iteration, train/eval mode, and a generic fit/evaluate loop.
  4. Add Conv2D, pooling, BatchNorm forward/backward behavior, registered running-statistic buffers, model summaries, state dictionaries, and safe save/load round trips.
  5. Port the MNIST MLP and train a CIFAR-10 CNN through the generic API; publish accuracy targets, API documentation, benchmarks, and limitations.

Hints in Layers

  1. Write the user-facing example before the internals; it becomes the API contract.
  2. Keep parameter discovery recursive and identity-based so shared parameters are not duplicated.
  3. Test serialization in a fresh process, not merely by reusing in-memory objects; verify BatchNorm output stays stable in evaluation mode after reload.

Common Pitfalls and Debugging

  • Optimizer updates miss some weights: parameters were stored in ordinary containers or not recursively registered. Test nested modules and lists explicitly.
  • Broadcasted operations return wrong gradients: backward must reduce expanded axes. Centralize an “unbroadcast” helper and test it independently.
  • Loaded model predicts differently: buffers, architecture metadata, or dtype were omitted. Version the state format and validate keys/shapes on load.
  • Evaluation changes with batch composition: BatchNorm is still using current-batch statistics. Propagate train/eval mode recursively and use saved running buffers during inference.

Definition of Done

  • Tensor/autograd operations pass unit and finite-difference tests, including broadcasting.
  • Nested/shared parameters are discovered exactly once and zeroed predictably.
  • SGD and Adam match hand-computed updates for several steps.
  • MLP and CNN examples train through one generic loop to declared accuracy targets.
  • BatchNorm gradients pass numerical checks, running statistics update only during training, and train/eval outputs match a trusted reference.
  • A CIFAR-10 CNN trains end to end through the public framework API and documents its reproducible validation accuracy.
  • Save/load in a fresh process preserves parameters and predictions.
  • Public API docs, model summary, tests, benchmarks, and known limitations are included.

Previous: Project 49 · Complete learning path · Next: Project 51