Models

Models#

Use this page to quickly reference how to define PyTorch models, inspect their inputs and outputs, examine trainable parameters, count parameters, and save or restore model state.

Defining Models#

Sequential#

Use Sequential when the computation is a simple ordered chain of modules.

from torch import nn

model = nn.Sequential(
    nn.Linear(4, 16),
    nn.ReLU(),
    nn.Linear(16, 3),
)

The output of each module becomes the input of the next one.

Module#

Subclass Module when a model needs custom structure or control flow.

from torch import nn

class Model(nn.Module):

    def __init__(self):
        super().__init__()
        self.layer = nn.Linear(4, 3)

    def forward(self, x):
        return self.layer(x)

Layers assigned as module attributes are registered automatically, including their parameters.

forward#

Define the forward computation inside forward.

def forward(self, x):
    hidden = self.hidden(x)
    hidden = torch.relu(hidden)
    return self.output(hidden)

Call the model itself:

outputs = model(inputs)

Warning

Call model(inputs), not model.forward(inputs). Calling the module lets PyTorch run the hooks and internal logic associated with nn.Module.

Common Layers#

Linear#

Applies an affine transformation to the last dimension of the input.

layer = torch.nn.Linear(in_features=128, out_features=10)

It transforms a tensor of shape (..., in_features) into a tensor of shape (..., out_features).

x = torch.randn(32, 128)

y = layer(x)  # → (32, 10)

With bias enabled, the parameter shapes are:

weight: (10, 128)
bias:   (10,)

For a layer with \(D_{\text{in}}\) input features and \(D_{\text{out}}\) output features, the parameter count is

\[ D_{\text{out}}D_{\text{in}} + D_{\text{out}}. \]

For example, the layer Linear(128, 10) has \(128 \times 10 + 10 = 1290\) parameters.

ReLU#

Applies the ReLU activation element-wise without changing the tensor shape.

activation = torch.nn.ReLU()
y = activation(x)

For one-off use, the functional form is also available:

y = torch.relu(x)

Flatten#

Flattens a range of dimensions inside a model.

flatten = torch.nn.Flatten(start_dim=1)

Using start_dim=1 preserves the batch dimension.

images = torch.randn(64, 1, 28, 28)

features = flatten(images)  # → (64, 784)

Dropout#

Behavior differs between training and evaluation.

  • During training, it randomly sets activations to zero and scales the remaining activations accordingly.

  • During evaluation, it does nothing.

dropout = torch.nn.Dropout(p=0.5)

The mode is controlled through model.train() and model.eval().

Identity#

Returns the input unchanged.

layer = torch.nn.Identity()

y = layer(x)  # y = x

Identity is useful when a branch or component should optionally perform no transformation.

Inputs & Outputs#

Batch Dimension#

Models normally operate on batches. The first dimension of an input tensor is conventionally designated as the batch dimension. The examples processed together in a batch are indexed along this dimension.

Single example:  (D,)
Tabular batch:   (N, D)

Single image:    (C, H, W)
Image batch:     (N, C, H, W)

Model Output#

The model output should preserve the batch dimension and follow the shape contract expected by the loss function. The output shape normally depends on the prediction task.

Task

Typical model output

Multiclass classification

logits (N, K)

Binary classification

logits (N,) or (N, 1)

Scalar regression

predictions (N,) or (N, 1)

Multi-output regression

predictions (N, T)

N = batch size – K = number of classes – T = number of predicted values per example.

See also

The required output and target contracts for different loss functions are covered in Training.

Logits#

Classification models normally return raw logits rather than probabilities.

For multiclass classification with K classes, a model returns K logits per example:

Input batch: (N, ...)
Logits:      (N, K)

For binary classification, a model can return one logit per example:

Input batch: (N, ...)
Logits:      (N,) or (N, 1)

A logit is an unrestricted real-valued score.

See also

Converting logits into probabilities or predicted classes is covered in Training and Evaluation.

Inspecting Parameters#

model.parameters()#

Iterate over all registered parameters recursively.

for parameter in model.parameters():
    print(parameter.shape)

This iterator is commonly passed directly to an optimizer:

optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

model.named_parameters()#

Inspect parameter names together with their tensors.

for name, parameter in model.named_parameters():
    print(name, parameter.shape)

This is useful when checking which part of a model owns a parameter.

requires_grad#

requires_grad controls whether autograd accumulates gradients for a parameter.

for parameter in model.parameters():
    print(parameter.requires_grad)

Disable gradient accumulation in place with:

parameter.requires_grad_(False)

To freeze all parameters in a model:

for parameter in model.parameters():
    parameter.requires_grad_(False)

The complete freezing and fine-tuning workflows are covered in Transfer Learning.

Total Parameter Count#

Count all trainable parameters in a model with:

trainable_parameters = 0
for parameter in model.parameters():
    if parameter.requires_grad:
        trainable_parameters += parameter.numel()

Count every parameter, including frozen ones, with:

total_parameters = sum(parameter.numel() for parameter in model.parameters())

Model State#

state_dict#

A model state dictionary maps parameter and persistent-buffer names to tensors.

state = model.state_dict()

Inspect its contents with:

for name, tensor in state.items():
    print(name, tensor.shape)

Buffers include stored values that are part of model state but are not trainable parameters, such as batch-normalization running statistics.

load_state_dict#

Restore model state with:

model.load_state_dict(state)

The model architecture must be compatible with the saved state.

Saving and Loading State#

Save a model state dictionary to disk:

torch.save(model.state_dict(), "model.pt")

Recreate the model architecture, load the state, and place the model on the required device:

model = Model()

state = torch.load("model.pt", map_location=device, weights_only=True)

model.load_state_dict(state)
model = model.to(device)

Warning

Saving a state dictionary does not save the Python code that defines the model architecture. Recreate a compatible model before calling load_state_dict.