---
title: BaseFinetuning
description: This class implements the base logic for writing your own Finetuning Callback.
canonical: >-
  https://lightning.ai/docs/pytorch/stable/optional-api/api/lightning.pytorch.callbacks.BaseFinetuning
---
# BaseFinetuning {#basefinetuning}

### `class lightning.pytorch.callbacks.BaseFinetuning` {#lightning.pytorch.callbacks.BaseFinetuning}

Bases: [`Callback`](/docs/pytorch/2.5.5/optional-api/api/lightning.pytorch.callbacks.Callback.md#lightning.pytorch.callbacks.Callback "lightning.pytorch.callbacks.callback.Callback")

This class implements the base logic for writing your own Finetuning Callback.

Warning

This is an [experimental](/docs/pytorch/2.5.5/reference/versioning.md#experimental-api) feature.

Override `freeze_before_training` and `finetune_function` methods with your own logic.

`freeze_before_training`: This method is called before `configure_optimizers`

and should be used to freeze any modules parameters.

`finetune_function`: This method is called on every train epoch start and should be used to

`unfreeze` any parameters. Those parameters need to be added in a new `param_group` within the optimizer.

Note

Make sure to filter the parameters based on `requires_grad`.

Example:

```python
>>> from torch.optim import Adam
>>> class MyModel(pl.LightningModule):
...     def configure_optimizer(self):
...         # Make sure to filter the parameters based on `requires_grad`
...         return Adam(filter(lambda p: p.requires_grad, self.parameters()))
...
>>> class FeatureExtractorFreezeUnfreeze(BaseFinetuning):
...     def __init__(self, unfreeze_at_epoch=10):
...         super().__init__()
...         self._unfreeze_at_epoch = unfreeze_at_epoch
...
...     def freeze_before_training(self, pl_module):
...         # freeze any module you want
...         # Here, we are freezing `feature_extractor`
...         self.freeze(pl_module.feature_extractor)
...
...     def finetune_function(self, pl_module, current_epoch, optimizer):
...         # When `current_epoch` is 10, feature_extractor will start training.
...         if current_epoch == self._unfreeze_at_epoch:
...             self.unfreeze_and_add_param_group(
...                 modules=pl_module.feature_extractor,
...                 optimizer=optimizer,
...                 train_bn=True,
...             )
```

#### `static filter_on_optimizer(optimizer, params)` {#lightning.pytorch.callbacks.BaseFinetuning.filter_on_optimizer}

This function is used to exclude any parameter which already exists in this optimizer.

Parameters:

*   **optimizer** ([`Optimizer`](https://docs.pytorch.org/docs/stable/optim.html#torch.optim.Optimizer "(in PyTorch v2.15)")) – Optimizer used for parameter exclusion
    
*   **params** ([`Iterable`](https://docs.python.org/3/library/collections.abc.html#collections.abc.Iterable "(in Python v3.15)")) – Iterable of parameters used to check against the provided optimizer
    

Return type:

[`list`](https://docs.python.org/3/builtins/stdtypes.html#list "(in Python v3.15)")

Returns:

List of parameters not contained in this optimizer param groups

#### `static filter_params(modules, train_bn=True, requires_grad=True)` {#lightning.pytorch.callbacks.BaseFinetuning.filter_params}

Yields the requires\_grad parameters of a given module or list of modules.

Parameters:

*   **modules** (`Union`\[[`Module`](https://docs.pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module "(in PyTorch v2.15)"), [`Iterable`](https://docs.python.org/3/library/collections.abc.html#collections.abc.Iterable "(in Python v3.15)")\[`Union`\[[`Module`](https://docs.pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module "(in PyTorch v2.15)"), [`Iterable`](https://docs.python.org/3/library/collections.abc.html#collections.abc.Iterable "(in Python v3.15)")\]\]\]) – A given module or an iterable of modules
    
*   **train\_bn** ([`bool`](https://docs.python.org/3/builtins/functions.html#bool "(in Python v3.15)")) – Whether not to train the BatchNorm module
    
*   **requires\_grad** ([`bool`](https://docs.python.org/3/builtins/functions.html#bool "(in Python v3.15)")) – Whether to create a generator for trainable or non-trainable parameters.
    

Return type:

[`Generator`](https://docs.python.org/3/library/collections.abc.html#collections.abc.Generator "(in Python v3.15)")

Returns:

Generator

#### `finetune_function(pl_module, epoch, optimizer)` {#lightning.pytorch.callbacks.BaseFinetuning.finetune_function}

Override to add your unfreeze logic.

Return type:

[`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")

#### `static flatten_modules(modules)` {#lightning.pytorch.callbacks.BaseFinetuning.flatten_modules}

This function is used to flatten a module or an iterable of modules into a list of its leaf modules (modules with no children) and parent modules that have parameters directly themselves.

Parameters:

**modules** (`Union`\[[`Module`](https://docs.pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module "(in PyTorch v2.15)"), [`Iterable`](https://docs.python.org/3/library/collections.abc.html#collections.abc.Iterable "(in Python v3.15)")\[`Union`\[[`Module`](https://docs.pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module "(in PyTorch v2.15)"), [`Iterable`](https://docs.python.org/3/library/collections.abc.html#collections.abc.Iterable "(in Python v3.15)")\]\]\]) – A given module or an iterable of modules

Return type:

[`list`](https://docs.python.org/3/builtins/stdtypes.html#list "(in Python v3.15)")\[[`Module`](https://docs.pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module "(in PyTorch v2.15)")\]

Returns:

List of modules

#### `static freeze(modules, train_bn=True)` {#lightning.pytorch.callbacks.BaseFinetuning.freeze}

Freezes the parameters of the provided modules.

Parameters:

*   **modules** (`Union`\[[`Module`](https://docs.pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module "(in PyTorch v2.15)"), [`Iterable`](https://docs.python.org/3/library/collections.abc.html#collections.abc.Iterable "(in Python v3.15)")\[`Union`\[[`Module`](https://docs.pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module "(in PyTorch v2.15)"), [`Iterable`](https://docs.python.org/3/library/collections.abc.html#collections.abc.Iterable "(in Python v3.15)")\]\]\]) – A given module or an iterable of modules
    
*   **train\_bn** ([`bool`](https://docs.python.org/3/builtins/functions.html#bool "(in Python v3.15)")) – If True, leave the BatchNorm layers in training mode
    

Return type:

[`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")

Returns:

None

#### `freeze_before_training(pl_module)` {#lightning.pytorch.callbacks.BaseFinetuning.freeze_before_training}

Override to add your freeze logic.

Return type:

[`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")

#### `static freeze_module(module)` {#lightning.pytorch.callbacks.BaseFinetuning.freeze_module}

Freezes the parameters of the provided module.

Parameters:

**module** ([`Module`](https://docs.pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module "(in PyTorch v2.15)")) – A given module

Return type:

[`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")

#### `load_state_dict(state_dict)` {#lightning.pytorch.callbacks.BaseFinetuning.load_state_dict}

Called when loading a checkpoint, implement to reload callback state given callback’s `state_dict`.

Parameters:

**state\_dict** ([`dict`](https://docs.python.org/3/builtins/stdtypes.html#dict "(in Python v3.15)")\[[`str`](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)"), [`Any`](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.15)")\]) – the callback state returned by `state_dict`.

Return type:

[`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")

#### `static make_trainable(modules)` {#lightning.pytorch.callbacks.BaseFinetuning.make_trainable}

Unfreezes the parameters of the provided modules.

Parameters:

**modules** (`Union`\[[`Module`](https://docs.pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module "(in PyTorch v2.15)"), [`Iterable`](https://docs.python.org/3/library/collections.abc.html#collections.abc.Iterable "(in Python v3.15)")\[`Union`\[[`Module`](https://docs.pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module "(in PyTorch v2.15)"), [`Iterable`](https://docs.python.org/3/library/collections.abc.html#collections.abc.Iterable "(in Python v3.15)")\]\]\]) – A given module or an iterable of modules

Return type:

[`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")

#### `on_fit_start(trainer, pl_module)` {#lightning.pytorch.callbacks.BaseFinetuning.on_fit_start}

Called when fit begins.

Return type:

[`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")

#### `on_train_epoch_start(trainer, pl_module)` {#lightning.pytorch.callbacks.BaseFinetuning.on_train_epoch_start}

Called when the epoch begins.

Return type:

[`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")

#### `setup(trainer, pl_module, stage)` {#lightning.pytorch.callbacks.BaseFinetuning.setup}

Called when fit, validate, test, predict, or tune begins.

Return type:

[`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")

#### `state_dict()` {#lightning.pytorch.callbacks.BaseFinetuning.state_dict}

Called when saving a checkpoint, implement to generate callback’s `state_dict`.

Return type:

[`dict`](https://docs.python.org/3/builtins/stdtypes.html#dict "(in Python v3.15)")\[[`str`](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)"), [`Any`](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.15)")\]

Returns:

A dictionary containing callback state.

#### `static unfreeze_and_add_param_group(modules, optimizer, lr=None, initial_denom_lr=10.0, train_bn=True)` {#lightning.pytorch.callbacks.BaseFinetuning.unfreeze_and_add_param_group}

Unfreezes a module and adds its parameters to an optimizer.

Parameters:

*   **modules** (`Union`\[[`Module`](https://docs.pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module "(in PyTorch v2.15)"), [`Iterable`](https://docs.python.org/3/library/collections.abc.html#collections.abc.Iterable "(in Python v3.15)")\[`Union`\[[`Module`](https://docs.pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module "(in PyTorch v2.15)"), [`Iterable`](https://docs.python.org/3/library/collections.abc.html#collections.abc.Iterable "(in Python v3.15)")\]\]\]) – A module or iterable of modules to unfreeze. Their parameters will be added to an optimizer as a new param group.
    
*   **optimizer** ([`Optimizer`](https://docs.pytorch.org/docs/stable/optim.html#torch.optim.Optimizer "(in PyTorch v2.15)")) – The provided optimizer will receive new parameters and will add them to add\_param\_group
    
*   **lr** ([`Optional`](https://docs.python.org/3/library/typing.html#typing.Optional "(in Python v3.15)")\[[`float`](https://docs.python.org/3/builtins/functions.html#float "(in Python v3.15)")\]) – Learning rate for the new param group.
    
*   **initial\_denom\_lr** ([`float`](https://docs.python.org/3/builtins/functions.html#float "(in Python v3.15)")) – If no lr is provided, the learning from the first param group will be used and divided by initial\_denom\_lr.
    
*   **train\_bn** ([`bool`](https://docs.python.org/3/builtins/functions.html#bool "(in Python v3.15)")) – Whether to train the BatchNormalization layers.
    

Return type:

[`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")

---

**Navigation**

- Parent: [Api](/docs/pytorch/2.5.5/optional-api/api.md)
