# Transfer Learning

Use this page to quickly reference how to load a pretrained model, replace its task-specific head, freeze or fine-tune parameters, and configure the optimizer for transfer learning.

## Pretrained Models

A pretrained model can be viewed as an encoder $f_{\theta}$ followed by a task-specific head $g_{\phi}$: 

$$
\mathbf{h}=f_{\theta}(\mathbf{x}),
\qquad
\hat{\mathbf{y}}=g_{\phi}(\mathbf{h}).
$$

The head maps the encoder representation to the output format expected by the original task. The encoder produces a representation that may generalize to new tasks. The idea of transfer learning is to reuse the pretrained encoder and replace the original head with another head whose output matches the new task.

### Loading a Pretrained Model

Torchvision models can be initialized from pretrained weights.

```python
from torchvision.models import resnet18, ResNet18_Weights

weights = ResNet18_Weights.DEFAULT
model = resnet18(weights=weights)
```

The pretrained weights define the input preprocessing expected by the model.

```python
transform = weights.transforms()
```

Use this as the canonical preprocessing for the pretrained model. See also [Data Management](quickref-data.md).

### Replacing the Head

A pretrained classifier ends with a task-specific head that maps the encoder representation to the classes used during pretraining. For a new task, the pretrained encoder can be reused, but the head must be replaced with new trainable layers that produce outputs compatible with the new task.

For example, the classification head of a pretrained ResNet is stored in `model.fc`.

```python
in_features = model.fc.in_features

model.fc = ... # Head for the new task
```

The location and name of the task-specific head depend on the model architecture.


## Training Conditions

During transfer learning, you must decide separately which parameters should receive gradients and which modules should use training or evaluation behaviour. These decisions depend on the transfer strategy, the characteristics of the pretrained model, and the new task.

### Parameters to Optimize

Parameters that are expected to be updated during training must belong to the optimizer and require a gradient. A common pattern is to optimize only parameters with the `requires_grad = True` flag.

```python
def get_trainable_parameters(model):
    for parameter in model.parameters():
        if parameter.requires_grad:
            yield parameter

optimizer = torch.optim.Adam(get_trainable_parameters(model), lr=1e-4)
```

Setting `requires_grad=False` normally prevents gradients from being accumulated for a parameter.

:::{important}
Changing `requires_grad` does not automatically change which parameters are already stored inside an existing optimizer. When the trainable parameter set changes, recreate the optimizer or deliberately reconfigure its parameter groups.
:::

### Parameter Groups

Optimizer parameter groups allow different parts of the model to use different settings.

```python
optimizer = torch.optim.Adam([
    {
        "params": model.layer4.parameters(),
        "lr": 1e-4,
    },
    {
        "params": model.fc.parameters(),
        "lr": 1e-3,
    },
])
```

### Frozen vs Eval

Freezing parameters and switching a module to evaluation mode control different things.

- `requires_grad=False` prevents gradient accumulation for a parameter.
- `module.eval()` changes the behaviour of modules such as dropout and batch normalization.

Calling `model.train()` or `model.eval()` does not change the `requires_grad` flag of any parameter.

:::{warning}
- Do not use `model.eval()` as a substitute for freezing parameters.
- Do not use `requires_grad=False` as a substitute for evaluation mode.
:::

### Batch Normalization

Freezing batch-normalization parameters does not freeze their running statistics. If the encoder remains in training mode, its batch-normalization layers can continue updating their stored statistics.

If a frozen encoder should behave entirely as a fixed feature extractor, place it in evaluation mode:

```python
encoder.eval()
head.train()
```

If the complete model is placed in training mode first, restore the encoder's evaluation mode afterward:

```python
model.train()
encoder.eval()
```

Calling `model.train()` recursively places its submodules in training mode, so the order matters.


## Transfer Workflows

Pretrained models can be adapted to a new task using different strategies. The appropriate strategy depends on the target data and task. Validation performance should guide the choice.


| Strategy                  | Encoder          | New head  |
| ------------------------- | ---------------- | --------- |
| Frozen feature extraction | Frozen           | Trainable |
| Partial fine-tuning       | Partly trainable | Trainable |
| Full fine-tuning          | Trainable        | Trainable |


### Reusable Transfer Model

A convenient way to enforce the training conditions of transfer learning is to wrap the encoder and head in a module that controls their modes and parameter freezing. The following example shows an implementation that can handle both frozen feature extraction and full fine-tuning.

```python
from torch import nn

class TransferModel(nn.Module):

    def __init__(self, encoder, head):
        super().__init__()
        self.encoder = encoder
        self.head = head
        self.encoder_frozen = False

    def forward(self, x):
        if self.encoder_frozen:
            with torch.no_grad():
                features = self.encoder(x)
        else:
            features = self.encoder(x)
        return self.head(features)

    def freeze_encoder(self):
        self.encoder_frozen = True
        for parameter in self.encoder.parameters():
            parameter.requires_grad_(False)
        self.encoder.eval()

    def unfreeze_encoder(self):
        self.encoder_frozen = False
        for parameter in self.encoder.parameters():
            parameter.requires_grad_(True)
        if self.training:
            self.encoder.train()

    def train(self, mode=True):
        super().train(mode)
        if mode and self.encoder_frozen:
            self.encoder.eval()
        return self
```

The encoder can be frozen or unfrozen explicitly. When it is frozen, the wrapper keeps it in evaluation mode during training; when it is unfrozen, it follows the mode of the complete model. When the encoder is frozen and gradients with respect to its inputs are not needed, its forward pass runs inside `torch.no_grad()`. Setting `requires_grad=False` prevents gradients from accumulating for the encoder parameters; `torch.no_grad()` additionally avoids recording the encoder computation for backpropagation.


### Frozen Feature Extraction

Frozen feature extraction reuses a pretrained encoder as a fixed feature extractor and trains only a new task-specific head. During training, the encoder stays in evaluation mode and its parameters do not receive gradients, while the new head remains trainable and operates in training mode.

| Component | `requires_grad` | Mode      |
| --------- | --------------: | --------- |
| Encoder   |         `False` | `eval()`  |
| New head  |          `True` | `train()` |

For example, a pretrained ResNet can be adapted to a new classification task by replacing its original head with an identity layer after exposing the encoder representation.

```python
from torchvision.models import resnet18, ResNet18_Weights

weights = ResNet18_Weights.DEFAULT
resnet = resnet18(weights=weights)

FEAT_DIM = resnet.fc.in_features
resnet.fc = nn.Identity()
```

A new head for the target task is created and combined with the frozen encoder.

```python
NUM_CLASSES = 10
head = nn.Linear(FEAT_DIM, NUM_CLASSES)

model = TransferModel(encoder=resnet, head=head)
model.freeze_encoder()

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

The resulting model is ready to be trained with the standard workflow explained in [Training](quickref-training.md). 
- Calling `model.train()` places the head in training mode while keeping the encoder in evaluation mode. 
- Calling `model.eval()` places both components in evaluation mode. 

:::{note}
Part of this example is specific to ResNet. Other model families organize their modules differently.
:::


### Full Fine-Tuning

Full fine-tuning allows the pretrained encoder to adapt to the new task, along with the new task-specific head. A common workflow first trains the new head with the encoder frozen, then unfreezes the encoder and continues training the complete model.

| Stage            | Encoder              | Head                 |
| ---------------- | -------------------- | -------------------- |
| Head training    | frozen, `eval()`     | trainable, `train()` |
| Full fine-tuning | trainable, `train()` | trainable, `train()` |

For example, a pretrained ResNet can be adapted to a new classification task by replacing its original head with an identity layer after exposing the encoder representation.

```python
weights = ResNet18_Weights.DEFAULT
resnet = resnet18(weights=weights)

FEAT_DIM = resnet.fc.in_features
resnet.fc = nn.Identity()

head = nn.Linear(FEAT_DIM, NUM_CLASSES)

model = TransferModel(encoder=resnet, head=head)
```

During the first stage, the encoder is frozen and only the new head is optimized. 

```python
model.freeze_encoder()

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

After the head has trained on the new task, the encoder is unfrozen and the optimizer is recreated so that the complete model can be updated. The learning rate is typically reduced for the second stage to avoid large updates to the pretrained parameters.

```python
model.unfreeze_encoder()

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

Part of this example is specific to ResNet. Other model families organize their modules differently.

:::{note}
Partial fine-tuning follows the same two-stage principle, but only selected encoder blocks are made trainable during the second stage. This requires architecture-specific control over both parameter freezing and module modes: trainable blocks operate in training mode, while frozen blocks containing stateful layers such as batch normalization remain in evaluation mode.
:::