5  Computation

Alongside giant datasets and powerful hardware, great software tools have played an indispensable role in the rapid progress of deep learning. Deep learning libraries let us recycle standard components while retaining the ability to modify anything, and over time their abstractions have grown coarser: from individual neurons, to layers, to the multi-layer blocks from which today’s models are assembled.

So far we called upon these libraries without asking how they work. We built models from layers, initialized and trained them, and treated everything between net(X) and the loss as machinery. This chapter opens the machinery. In 2016 that meant learning to build, initialize, and save a small model trained from scratch in 32-bit arithmetic on one device. Those skills remain, and this chapter teaches them, but the working assumptions around them have changed: a model is now assembled from a configuration object, measured in gigabytes, run in reduced precision, checkpointed together with its optimizer state, and as often as not initialized from someone else’s weights rather than from a random number generator.

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 to write layers the library does not provide (Section 5.4). We then turn to the numeric formats models compute in (Section 5.5), how state is saved, restored, and adopted from pretrained models (Section 5.6), how tensors and models live on GPUs and in GPU memory (Section 5.7), and finally how to make runs repeatable and inspect a model from the outside (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 go deeper on the machinery this chapter opens up: 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 walks tensors, storage, autograd, and nn.Module mechanics at exactly this chapter’s level of “open the machinery”.
• 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; you build “needle”, a miniature framework with autograd, modules, initialization, and GPU support — this entire chapter from the implementor’s side.
• Neural Networks: Zero to Hero — Andrej Karpathy — free video series; builds autograd, modules, and training loops from scratch in plain Python, making every abstraction in this chapter concrete before you rely on the library’s version.
• fast.ai Part 2: Deep Learning Foundations — Howard et al. — free; rebuilds a training framework from tensor ops upward (modules, initialization, mixed precision, accelerated training), the practitioner’s companion to this chapter.

Tutorials, notes, and interactive

• PyTorch internals — Edward Yang — free; the classic 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; the compute-, memory-, and overhead-bound mental model that Section 5.7 builds on.
• Reproducibility — PyTorch notes — free; the determinism flags, their costs, and their limits — the fine print behind 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.