# 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.

```python
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.

```python
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`.

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

Call the model itself:

```python
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.

```python
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)`.

```python
x = torch.randn(32, 128)

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

With bias enabled, the parameter shapes are:

```text
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.

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

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

```python
y = torch.relu(x)
```


### `Flatten`

Flattens a range of dimensions inside a model.

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

Using `start_dim=1` preserves the batch dimension.

```python
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.

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

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


### `Identity`

Returns the input unchanged.

```python
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.

```text
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.

:::{seealso}
The required output and target contracts for different loss functions are covered in [Training](quickref-training.md).
:::

### Logits

Classification models normally return raw logits rather than probabilities.

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

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

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

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

A logit is an unrestricted real-valued score. 

:::{seealso}
Converting logits into probabilities or predicted classes is covered in [Training](quickref-training.md) and [Evaluation](quickref-evaluation.md). 
:::

## Inspecting Parameters

### `model.parameters()`

Iterate over all registered parameters recursively.

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

This iterator is commonly passed directly to an optimizer:

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


### `model.named_parameters()`

Inspect parameter names together with their tensors.

```python
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.

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

Disable gradient accumulation in place with:

```python
parameter.requires_grad_(False)
```

To freeze all parameters in a model:

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

The complete freezing and fine-tuning workflows are covered in [Transfer Learning](quickref-transfer.md).

### Total Parameter Count

Count all trainable parameters in a model with:

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

Count every parameter, including frozen ones, with:

```python
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.

```python
state = model.state_dict()
```

Inspect its contents with:

```python
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:

```python
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:

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

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

```python
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`.
:::