Skip to content

Module & Layer Engineering

Overview

An nn.Module is more than a bag of layers: it's a declarative container whose tree structure drives device placement, serialization, and gradient wiring. Mastery means knowing exactly how the tree is discovered, how to inject logic (hooks), and how to build composable custom layers.

Topics

  • Custom Layers & Advanced Containers - Writing nn.Module subclasses with shared primitive ops, ModuleList/ModuleDict, and init rules.
  • Weight Initialization - Why zero init kills nets, Xavier/Kaiming math, and how reset_parameters is wired.
  • Parameter Sharing & Weight Tying - One storage, many roles; tying encoder/decoder or embedding/output matrices.
  • Hooks — Debug, Inject, Extract - Forward/backward hooks to snoop activations, patch behavior, and inspect grads without editing forward.

Key Patterns

  • Always call super().__init__(); the module tree is built by attribute assignment.
  • Assign submodules / registered params as attributes — never in a Python list.
  • Initialize in __init__ (or reset_parameters) rather than lazily in forward.
  • Hooks are the non-invasive way to instrument, not subclassing.