🎓 A1Learner
🔍

PyTorch What Is a Module

A module is a building block that holds parts of a model and runs maths on input. In PyTorch, modules come from the base class nn.Module.

A module is a class

The class extends nn.Module. __init__ builds the parts, and forward does the maths.

class Doubler(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        return x * 2

Call the model

Calling the model runs forward on the input. Writing model(numbers) is the normal way to run a module.

model = Doubler()
numbers = torch.tensor([1.0, 3.0, -4.0])
print(model(numbers))
print(model.forward(numbers))

A module is an object

A built module is a Python object with a type. The type name is the class name.

print(type(model).__name__)
print(isinstance(model, nn.Module))

Example

import torch
from torch import nn


class Doubler(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        return x * 2


model = Doubler()
numbers = torch.tensor([1.0, 3.0, -4.0])
result = model(numbers)

print("Input:", numbers)
print("Output:", result)
print("Is a module:", isinstance(model, nn.Module))

Key points

  • Module: a building block that can hold layers and run maths
  • Use class Name(nn.Module) to define a new module
  • __init__: the method that builds and stores the parts
  • forward: the method that receives input and returns output
  • Use model(x) to run a module on input x
  • A module bundles its parts and its maths in one class

In this example, the Doubler class extends nn.Module. Its __init__ calls super().__init__() and stores no parts. Its forward returns x * 2. The program builds one Doubler and one tensor of three numbers. Calling model(numbers) runs forward on the tensor. The output doubles each number, minus signs included. The last line confirms the model is an nn.Module.

Note

A module can hold other modules, so small blocks build big models

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

What does the forward method of a module do?