6.8 PyTorch: tensors, autograd, modules

What this is and why it exists

PyTorch is the common language of research and much of industry, and fluency in it means that any architecture you can sketch on paper, you can build. There are only three ideas to hold: an array type that knows where it lives, a system that records operations so it can differentiate them, and a container that organises parameters into a model. Almost every error a beginner meets is a mismatch in the first, or a misunderstanding of when the second is recording.

The vocabulary

• Tensor — the array type, with a shape, a data type, and a device.
• Device — where the tensor's memory lives: the processor or an accelerator.
• dtype — what the numbers are, which decides precision and memory.
• autograd — the system that records operations so gradients can be computed.
• requires_grad — the flag marking a tensor as something to differentiate with respect to.
• Module — the base class every model component inherits from.
• Parameter — a tensor the optimiser will update; buffer — one it will not.
• state_dict — the dictionary of a module's parameters and buffers.

The mental model

A tensor has three properties and most beginner errors are one of them. Its shape, which decides whether an operation is defined at all. Its data type, which decides precision and whether an operation is even permitted — an integer tensor where a float is expected is a common one. And its device, because an accelerator cannot read the processor's memory: an operation between a tensor on one and a tensor on the other fails, and the fix is .to(device) on whichever is in the wrong place. Write the device once at the top of a script, move the model and every batch to it, and the whole class of error disappears.

Two shape behaviours are worth naming because they cause silent bugs rather than crashes. Broadcasting expands a smaller shape to meet a larger one, which is convenient and will happily make an operation succeed that you did not intend — a prediction of one shape against a target of another can broadcast into a nonsense loss that trains to nothing. And the difference between a view and a copy decides whether modifying one array changes another. When something is wrong and nothing has crashed, print shapes.

Autograd records as you go. PyTorch's own description: it "records a graph recording all of the operations that created the data as you execute operations, giving you a directed acyclic graph whose leaves are the input tensors and roots are the output tensors", and "the graph is recreated from scratch at every iteration, and this is exactly what allows for using arbitrary Python control flow statements". The flag that starts it is requires_grad, "a flag, defaulting to false unless wrapped in a nn.Parameter, that allows for fine-grained exclusion of subgraphs from gradient computation", and the rule for what gets recorded is that "an operation is only recorded in the backward graph if at least one of its input tensors require grad."

Which brings us to the mistake that costs memory rather than correctness. In an evaluation loop you do not want gradients, but the model's parameters still require them, so every operation is recorded and every intermediate value is retained for a backward pass that never comes. Memory climbs until the run stops. The context manager is the fix, and its behaviour is documented exactly: "computations in no-grad mode behave as if none of the inputs require grad. In other words, computations in no-grad mode are never recorded in the backward graph even if there are inputs that have require_grad=True."

detach() is the other exit and it does something different: it gives you the same values as a tensor outside the graph, which is what you want when you accumulate a running loss for logging. Adding a loss tensor itself to a total keeps the entire graph of every batch alive; the symptom is memory growing steadily through an epoch, and the fix is to record a plain number rather than a live tensor.

A module is a container that knows what it owns. Subclass it, register child modules and parameters as attributes, define what a forward pass does, and you get the parameter collection, the device movement and the mode switching for free. The distinction that matters is between a parameter, which the optimiser updates, and a buffer, which travels with the model and is saved and moved but never trained — a normalisation layer's running estimates are the standard example. Register a buffer rather than keeping a bare tensor as an attribute, or it will not move to the accelerator with the model and will not be saved.

Saving means saving the state dictionary, not the object. The PyTorch documentation is explicit: "Instead of saving a module directly, for compatibility reasons it is recommended to instead save only its state dict", and "saving a module's state_dict is a best practice when using torch.save." Saving a whole object pickles your class definitions along with it, so the file only loads where that exact code exists, and it breaks on refactoring. Saving the dictionary means the file is data — reconstruct the model in code, load the dictionary into it, and the checkpoint survives.

One security point belongs here, because a checkpoint is a file you may have obtained from elsewhere. Loading with the weights-only option restricts what the file is permitted to reconstruct; the documentation notes it "narrows the surface of remote code execution attacks" while also being clear about the limits — it "does not guard against denial of service attacks", and an old checkpoint containing a module object needs the option turned off to load at all. Treat a checkpoint from an untrusted source the way you would treat a program from an untrusted source, because without that option it effectively is one.

Two further habits. Save the optimiser's state alongside the model, because resuming without it restarts the momentum and adaptive estimates from nothing and the first steps after a resume are wrong. And save the epoch number and the configuration, so a checkpoint answers what it is rather than only what it weighs.

What you should now be able to explain or do

Name the three properties of a tensor and say which error each produces. Recognise broadcasting and view-versus-copy as sources of silent bugs. Say what autograd records and what starts the recording. Use no-grad in an evaluation loop and explain what changes. Use detach for logging and explain the memory symptom it fixes. Build a module, and distinguish a parameter from a buffer. Save and load a state dictionary, with the optimiser state, and say why saving the object is worse. State the security consideration when loading a checkpoint you did not produce.

Check yourself