torch.nn.Module packages a model’s trainable state and computation into one reusable object. Define Parameters in __init__(), implement the forward pass in forward(), and call the instance on inputs. section 3.4.2
A minimal linear model
class LinearModel(torch.nn.Module): def __init__(self): super().__init__() self.W = torch.nn.Parameter(torch.rand(input_dim, output_dim)) self.b = torch.nn.Parameter(torch.zeros(output_dim)) def forward(self, inputs): return torch.matmul(inputs, self.W) + self.b Instantiate it with model = LinearModel(), then use output = model(torch_inputs). Calling model(inputs), rather than model.forward(inputs), preserves the framework hooks added by __call__(). section 3.4.2
A Module simplifies the training loop by handing state and updates to PyTorch components:
- Retrieve trainable Parameters with model.parameters() and pass them to torch.optim.SGD. section 3.4.2
- Compute predictions and loss, call loss.backward(), update with optimizer.step(), and clear gradients with model.zero_grad(). section 3.4.2
- Use Modules as building blocks within modular, hierarchical architectures and reuse them where appropriate. section 9.1
Listing 3.27 Defining a torch.nn.Module
class LinearModel(torch.nn.Module): def __init__(self): super().__init__() self.W = torch.nn.Parameter(torch.rand(input_dim, output_dim)) self.b = torch.nn.Parameter(torch.zeros(output_dim)) def forward(self, inputs): return torch.matmul(inputs, self.W) + self.b
Inspect the complete minimal Module definition, including superclass initialization, Parameter fields, and the forward method. section 3.4.2
