Unit 2.3 · about 4 h
179 / 324Modules, parameters, buffers and initialization
Parameters, buffers and init are the state you own for the rest of the stage. Confusing them turns every later resume, transfer and device bug into a mystery.
- explain
- implement
You can now
- Assemble nested modules and read ownership from dotted names in
named_parameters. - Separate Parameters, persistent buffers and non-persistent buffers, and predict exactly which names appear in
state_dict. - Replace default initialization explicitly, and reload saved state into a fresh instance.
Covers
- nn.Module, Parameter, buffers and nested registration
- Functional versus module APIs
- Explicit initialization and state_dict semantics
Study, in this order
3 steps
- Build the Neural Network: class, layers, parametersPyTorch - Learn the Basics
Define the Class (
nn.Modulesubclass,forward, never callmodel.forward()directly); Model Layers (Flatten,Linear,ReLU,Sequential,Softmax); Model Parameters (named_parameterswith dotted names) - Modules note: nesting, state, buffers, initializationPyTorch - Autograd mechanics, serialization and reproducibility
A Simple Custom Module; Modules as Building Blocks (children, recursion,
ModuleListandModuleDict,to(),apply()with xavier init); Module State (parameters, persistent and non-persistent buffers,StatefulModulekey list); Module Initialization (skip_initwithorthogonal_)Depth · Skip Module Hooks and Advanced Features; they belong to later profiling and distributed work.
- Extending `torch.nn`: the module recommendation and the `Linear` examplePyTorch - Autograd mechanics, serialization and reproducibility
Extending
torch.nnintro (modules for stateful layers, functional form for parameter-less ops); Adding a Module (nn.Parameter,register_parameter,register_buffer,nn.init.uniform_comments)Depth · Read only the Adding a Module part here; the autograd Function part belongs to the custom-autograd unit.
Focus on
- Assigning
nn.Parameterregisters; assigning a plainTensordoes not. Registration is whatparameters(),.to()and optimizers see. - Persistent buffers ship in
state_dict, non-persistent ones do not, but.to()moves both. state_dictmaps dotted names to tensors only. Reloading needs a matching live instance first; mismatched names fail unless explicitly allowed.- Default init runs at construction. To study init, skip it with
skip_initand apply onenn.initfunction yourself, as the Modules note does withorthogonal_.
Reproduce closed-book
Linear submodule, a persistent running-mean buffer and a ParameterList. Predict the state_dict keys, save, reload into a fresh instance, and assert the buffer matches.Traps
- Expecting a plain-Tensor attribute to train. Only Parameters and their submodules appear in
parameters(). - Saving the model object instead of its
state_dict. Pickle then binds the file to the class definition;state_dictplus code reloads cleanly. - Copying the Extending example
uniform(-0.1, 0.1)init into real work. It ignores fan-in; use a fan-awarenn.initfunction orskip_initplus your own.
Practice
Learn PyTorch 01: nn.Module with Parameters, then save and reload (tasks 2 and 5)
Exercises, task 2 (nn.Module subclass with nn.Parameter weight and bias, state_dict check) and task 5 (save state_dict, reload into a new instance, verify predictions match)
Implement in your own script. The prediction check must compare outputs, not just file existence.
Self-test
-
Which names appear in
StatefulModulestate_dict, and which entries are absent?param1,param2,buffer1, theparam_listandparam_dictentries,linear.weightandlinear.biasappear. The non-persistentbuffer2and the reservedNoneentriesparam3andbuffer3are absent. -
Why does
load_state_dictneed a constructed model first?The dict holds tensors under names, not the module graph. The instance supplies the structure the names attach to. -
A
BatchNormlayer updates statistics in train mode and freezes them in eval mode. Where do those statistics live?In persistent buffers. They are saved instate_dictand moved by.to(), but never given to the optimizer. -
What does
module.apply(init_weights)give you that a loop overnamed_parametersdoes not?apply calls the function on every submodule and then on the module itself, so one function with anisinstancecheck can initialize each layer type its own way. Anamed_parametersloop sees only names and tensors, not the layer each one belongs to.
Completion criterion
Stop when
Your predicted key list matches the printed state_dict before you print it, and the reloaded instance reproduces predictions exactly.
Leads to
- 2.4 Backward, `zero_grad`, `gradcheck` and evaluation modes Drive the parameters you can now take apart.
- 2.7 Explicit training and evaluation loops with correct metrics Train the module you can now inspect.
- 3.1 Initialization and signal scale across depth Go deeper on why init scales matter.
Comes after
- 1.7 Aliasing, mutation and module ownershipMeet parameters, buffers and init in the framework.
- 2.2 Autograd state: views, in-place ops, leaves and accumulationParameters are the leaves you must stop detaching by accident.
Resources for this unit
- PyTorch - Learn the Basics Complete Tensors, Datasets and DataLoaders, Transforms, Build Model, Autograd, Optimization, and Save and Load Model. Rebuild each essential example in a small src package instead of retaining a tutorial notebook.
- PyTorch - Autograd mechanics, serialization and reproducibility Read Autograd mechanics, Serialization semantics and Reproducibility. Test leaf behavior, saved tensors, grad modes, state dictionaries and the exact tolerance promised by one interrupted and resumed run.
Also in this section
- Learn PyTorch - selected framework exercises Complete exercise 00, tasks 1-5 in 01, tasks 1-7 in 02, tasks 5-13 in 03, tasks 1-3 in 05, and exercises 1-2 in 07. Recreate solutions independently and verify behavior against current official docs.
- Deep Learning with PyTorch, 2nd edition - framework foundations Read Sections 3.8-3.9 and 5.5 when tensor storage, strides or autograd need a second explanation, with the Chapter 3 and 5 companion notebooks for runnable cells. Inspect the Chapter 4 and 7 code only for a diagnosed data-representation or loading gap.
- NeetCode ML - PyTorch closing drills After the module gate, solve Pytorch Basics, Layer Normalization, Batch Normalization and RMS Normalization cold. Explain shapes, state and train-evaluation behavior aloud. These do not replace the package and resume test.
- PyTorch Automatic Mixed Precision recipe - Saving/Resuming Read the Saving/Resuming section: save
scaler.state_dict()with the model and optimizer state, at the start of an iteration or afterscaler.update(), and how to resume with or without AMP. - Hands-On Machine Learning with Scikit-Learn and PyTorch Read Chapter 10, Implementing Mini-Batch Gradient Descent Using DataLoaders and Model Evaluation, for the train and evaluate loops and how batch metrics must be aggregated.