5  Computation

Deep learning libraries provide reusable layers and training components while allowing models to define their own computation. Their central abstraction is the module: a unit that can contain parameters, child modules, and a forward computation. Modules range from individual layers to the repeated blocks from which large models are assembled.

This chapter explains the software structure behind the models used so far. It covers module construction, initialization, model state, serialization, custom layers, reproducibility, numeric formats, and device placement. These topics matter at any scale, but modern training makes their interaction especially important: models are commonly built from configuration objects, run in reduced precision, measured in gigabytes, initialized from pretrained weights, and checkpointed together with optimizer state.

Accordingly, we proceed in eight steps. We start with how models are built from modules and configs (Section 5.1), what a model’s state is and what it costs in memory (Section 5.2), how that state is initialized (Section 5.3), and how that state is saved, restored, and adopted from pretrained models (Section 5.6). We then turn to numeric formats (Section 5.5), devices and GPU memory (Section 5.7), and layers the library does not provide (Section 5.4). Finally, we distinguish repeatable experiments from inspecting a model’s execution (Section 5.8). The chapter introduces no new models or datasets; the advanced modeling chapters that follow rely on these techniques throughout.

Resources and Further Reading

The references below develop the mechanisms introduced in this chapter: how frameworks represent models and state, automatic differentiation, numerics and mixed precision, devices and memory, and reproducible training. All are freely accessible online except where noted.

Books

• Deep Learning with PyTorch — Stevens, Antiga & Viehmann — free PDF from the PyTorch team (archived copy; the original pytorch.org link has gone away); Part 1 covers tensors, storage, autograd, and nn.Module at the same level as this chapter.
• Machine Learning Systems — Vijay Janapa Reddi — free online; the systems view around this chapter: frameworks, data pipelines, training infrastructure, and efficient deployment.

Courses and video lectures

• CMU 10-414/714: Deep Learning Systems — Chen & Kolter — free lectures and assignments; students build “needle”, a miniature framework with autograd, modules, initialization, and GPU support, providing an implementer’s view of these abstractions.
• Neural Networks: Zero to Hero — Andrej Karpathy — free video series; builds autograd, modules, and training loops from scratch in plain Python, illustrating the abstractions used by deep learning libraries.
• fast.ai Part 2: Deep Learning Foundations — Howard et al. — free; rebuilds a training framework from tensor operations upward, including modules, initialization, mixed precision, and accelerated training.

Tutorials, notes, and interactive

• PyTorch internals — Edward Yang — free; a guided tour of tensors, strides, dispatch, and autograd inside PyTorch, one level below Section 5.1.
• JAX — The Sharp Bits — free; pure functions, explicit PRNG keys, and jit constraints — the functional worldview behind this book’s JAX tab, stated as a list of gotchas.
• Train With Mixed Precision — NVIDIA — free; the vendor guide to fp16/bf16 arithmetic and loss scaling that Section 5.5 distills.
• What Every Computer Scientist Should Know About Floating-Point Arithmetic — David Goldberg — free; the standard reference beneath every rounding and overflow issue in Section 5.5.
• Making Deep Learning Go Brrrr From First Principles — Horace He — free; a first-principles account of compute, memory, and framework overhead covered in Section 5.7.
• Reproducibility — PyTorch notes — free; the determinism flags, their costs, and their limits, complementing Section 5.8.

Foundational papers

• Automatic Differentiation in Machine Learning: a Survey — Baydin, Pearlmutter, Radul & Siskind (2018), JMLR — free; the definitive account of forward- and reverse-mode autodiff, the algorithm every framework in this book implements.
• Mixed Precision Training — Micikevicius et al. (2018) — free; the origin of the fp16 + master-weights + loss-scaling recipe in Section 5.5.
• PyTorch: An Imperative Style, High-Performance Deep Learning Library — Paszke et al. (2019) — free; the design rationale (eager execution, autograd, memory allocator) for the imperative style this chapter teaches.