Fine-tuning

ModelFreezer

class pytorch_accelerated.finetuning.ModelFreezer(model, freeze_batch_norms=False)

A class to freeze and unfreeze different parts of a model, to simplify the process of fine-tuning during transfer learning.

This class uses the following abstractions:

• Layer: A subclass of torch.nn.Module with a depth of 1. i.e. The module is not nested.

• LayerGroup: The modules which have been defined as attributes of a model. These may be Layers or nested modules.

For example, let’s consider the following model:

from torch import nn

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.input = nn.Linear(100, 100)
        self.block_1 = nn.Sequential(
            nn.Linear(100, 100),
            nn.BatchNorm1d(100),
            nn.Sequential(
                nn.Linear(100, 100),
                nn.BatchNorm1d(100),
                nn.ReLU(),
            ),
        )
        self.output = nn.Linear(100, 10)

    def forward(self, x):
        x = self.input(x)
        x = self.block_1(x)
        out = self.output(x)
        return out

Here, the layer groups would be the modules [input, block_1, output], whereas the layers would be ordered, flattened list of Linear, BatchNorm and ReLU modules.

freeze(from_index=0, to_index=-2, set_modules_as_eval=False)

Freeze layer groups corresponding to the specified indexes, which are inclusive. By default, this freezes all layer groups except the final one.

Parameters:

• from_index – The index of the first layer group to freeze.

• to_index – The index of the final layer group to freeze.

• set_modules_as_eval – If True, frozen modules will also be placed in eval mode. This is False by default.

get_layer_groups() → List[LayerGroup]

Return all of the model’s layer groups. A layer group is any module which has been defined as an attribute of the model.

Returns:

a list of all layer groups in the model.

get_layers() → List[Layer]

Return all of the model’s layers. A Layer is any non-nested module which is included in the model.

Returns:

a list of all layers in the model.

get_trainable_parameters()

Return a list of all unfrozen model parameters, which will be updated during training.

Returns:

a list of all trainable parameters

unfreeze(from_index=-1, to_index=0, set_modules_as_training=True)

Unfreeze layer groups corresponding to the specified indexes, which are inclusive. By default, this unfreezes all layer groups. For each layer group, any parameters which have been unfrozen are returned, so that they can be added to an optimizer if needed.

Parameters:

• from_index – The index of the first layer group to unfreeze.

• to_index – The index of the final layer group to unfreeze.

• set_modules_as_training – If True, unfrozen modules will also be placed in train mode. This is True by default.

Returns:

a dictionary containing the parameters which have been unfrozen for each layer group.