Gradually unfreezing layers: UserWarning: The provided params to be frozen already exist within another group of this optimizer. Those parameters will be skipped

Hello,

I am trying to extend the pytorch lightning class pytorch_lightning.callbacks.BaseFinetuning to unfreeze the layers of my network gradually.

My network is composed of 3 groups of modules:

  • Features
  • Middle layers
  • Outputs

Here, I am trying to gradually unfreeze the features layers after a certain epoch. By gradually, I mean that I want to unfreeze one layer after every epoch.

Here’s my class:

from typing import Generator, Iterable, Optional, Union

from pytorch_lightning.callbacks import BaseFinetuning
import torch.nn
from torch.nn import Module
from torch.nn.modules.batchnorm import _BatchNorm


class CustomFinetuning(BaseFinetuning):

    def __init__(
            self,
            unfreeze_at_epoch: Optional[int] = None,
            refreeze_at_epoch: Optional[int] = None,
            gradual_unfreeze: bool = False
    ):
        super().__init__()
        self._unfreeze_at_epoch = unfreeze_at_epoch
        self._refreeze_at_epoch = refreeze_at_epoch
        self._gradual_unfreeze = gradual_unfreeze

    def freeze_before_training(self, pl_module):
        # freeze any module you want
        # Here, we are freezing `features`
        self.freeze(pl_module.net.features)

    def finetune_function(self, pl_module, current_epoch, optimizer, optimizer_idx):
        # When `current_epoch` is self._unfreeze_at_epoch, features layers will start training.
        if self._unfreeze_at_epoch is not None:
            if self._gradual_unfreeze is False and current_epoch == self._unfreeze_at_epoch:
                self.unfreeze_and_add_param_group(
                    modules=pl_module.net.features,
                    optimizer=optimizer,
                    train_bn=True,
                )
            if self._gradual_unfreeze is True and current_epoch >= self._unfreeze_at_epoch:
                modules_to_unfreeze = list(self.filter_modules(pl_module.net.features, train_bn=False, requires_grad=False))
                if len(modules_to_unfreeze) > 0:
                    self.unfreeze_and_add_param_group(
                        modules=[modules_to_unfreeze[-1]],
                        optimizer=optimizer,
                        train_bn=True
                    )
        if self._refreeze_at_epoch is not None:
            if current_epoch == self._refreeze_at_epoch:
                self.freeze(pl_module.net.features)
                self.freeze(pl_module.net.middle_layers)

    @staticmethod
    def is_frozen(module: torch.nn.Module):
        return all([
            param.requires_grad is False for param in module.parameters()
        ])

    @staticmethod
    def filter_modules(
        modules: Union[Module, Iterable[Union[Module, Iterable]]], train_bn: bool = True, requires_grad: bool = True
    ) -> Generator:
        """Yields the modules with the `requires_grad` parameters of a given module or list of modules.

        Args:
            modules: A given module or an iterable of modules
            train_bn: Whether to train BatchNorm module
            requires_grad: Whether to create a generator for trainable or non-trainable parameters.
        Returns:
            Generator
        """
        modules = BaseFinetuning.flatten_modules(modules)
        for mod in modules:
            if isinstance(mod, _BatchNorm) and not train_bn:
                continue
            # recursion could yield duplicate parameters for parent modules w/ parameters so disabling it
            if any([
                param.requires_grad == requires_grad
                for param in mod.parameters(recurse=False)
            ]):
                yield mod

The code works fine when gradual_unfreeze=False, that is, when all the features layers are unfrozen at the same time.
However, the gradual_unfreeze=True, I get the following warning when calling unfreeze_and_add_param_group:

/Users/matteo/Documents/Fusus/Coding/ai-models/venv/lib/python3.8/site-packages/pytorch_lightning/callbacks/finetuning.py:211: UserWarning: The provided params to be frozen already exist within another group of this optimizer. Those parameters will be skipped.
HINT: Did you init your optimizer in `configure_optimizer` as such:
 <class 'torch.optim.adamw.AdamW'>(filter(lambda p: p.requires_grad, self.parameters()), ...) 
  rank_zero_warn(

The optimizer is simply initialized this way:

optimizer = self.hparams.optimizer(params=self.parameters())

The warning is raise at the first tie the function unfreeze_and_add_param_group is called. At that time, modules_to_unfreeze[-1] is a Conv2d(320, 1280, kernel_size=(1, 1), stride=(1, 1), bias=False).

I’d like to get rid of the warning so that I can be sure of what is actually being unfrozen. Do you know what could be the issue? Is there any other information I can provide you to help?

Thanks

Based on this code you are already passing all parameters to the optimizer, i.e. the frozen and trainable ones.
Afterwards it seems you are unfreezing some parameters and try to pass them again to the same optimizer via add_param_group, which will fail since these parameters were already passed to the optimizer. Either filter these frozen parameters out initially as explained in the error message or just unfreeze them without trying to create a new parameter group.

@ptrblck , thank you for your reply. You are right. I also checked the torch lightning code and it looks like the warning is not an actual in my case.