Builders' Guide

The last three chapters called on the library and skipped over how it works. Now we peel back the curtain. None of this introduces a new kind of network; all of it is the plumbing the advanced chapters lean on constantly — and it turns out a surprising amount of that plumbing is worth watching run.

The module is the one abstraction that scales: a layer is a module, a block of layers is a module, and the whole model is a module. Because they nest, a handful of lines can describe a network hundreds of layers deep. And because a forward is just Python, a module can do anything a function can.

That flexibility explains a convenience you've been using without noticing: you never once specified an input dimension. The framework waits for data, then fills in every shape at once.

Once the shapes exist, the parameters have to start somewhere. The default is usually fine, but when you move off the beaten path — a custom block, a reproducible experiment, an idea from a paper — you'll want to set the initialization yourself.

Parameters aren't only for training. You reach into them to debug, to share a representation between two parts of a model, and — most of all — to save your work.

The lab

Three pieces of the machinery, yours to drive: compose a block and watch its shapes resolve, roll your own weight initializer, and meet a layer that has no parameters at all.

Using a GPU

The chapter ends where real training begins: hardware. GPU throughput has grown roughly a thousand-fold per decade, and harnessing it is mostly about where your data lives.

The real code

Every act above carries a </> chip with its section's exact PyTorch. The recurring shape is a Module subclass with a constructor and a forward — everything else (parameters, backprop, serialization) the base class handles:

class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.hidden = nn.LazyLinear(256)
        self.out = nn.LazyLinear(10)
 
    def forward(self, X):
        return self.out(F.relu(self.hidden(X)))

# a parameter-free custom layer
class CenteredLayer(nn.Module):
    def forward(self, X):
        return X - X.mean()
 
# a layer with its own parameters
class MyLinear(nn.Module):
    def __init__(self, in_units, units):
        super().__init__()
        self.weight = nn.Parameter(torch.randn(in_units, units))
        self.bias = nn.Parameter(torch.randn(units,))
    def forward(self, X):
        return F.relu(torch.matmul(X, self.weight.data) + self.bias.data)
 
# tied parameters: one shared layer, used twice
shared = nn.LazyLinear(8)
net = nn.Sequential(nn.LazyLinear(8), nn.ReLU(), shared, nn.ReLU(),
                    shared, nn.ReLU(), nn.LazyLinear(1))
 
# save / load the parameter dictionary
torch.save(net.state_dict(), 'mlp.params')
clone = MLP()
clone.load_state_dict(torch.load('mlp.params'))

Three things to remember

1. Everything is a module. A layer, a block, and the whole model share one interface — ingest inputs, return outputs, hold parameters — and because modules nest, you compose enormous networks from tiny, reusable parts. A forward can run arbitrary code, loops and all.
2. Shapes are lazy; parameters are explicit. The framework infers every weight shape the first time data flows through, which is why you never specify an input dimension — but how those weights are initialized is a choice with real consequences, from dead symmetric networks to variance-preserving Xavier.
3. Weights are data; the architecture is code. You save a model by saving its parameter dictionary and rebuilding the network in code — the same reason two layers can share one tensor, and the same reason moving that data to a GPU is something you do on purpose.

Adapted from Chapter 6 of Dive into Deep Learning by Zhang, Lipton, Li, and Smola (CC BY-SA 4.0). The figures run this site's verified TypeScript port; all computation happens live in your browser.