def unfreeze_and_add_param_group(module: torch.nn.Module,
optimizer: Optimizer,
lr: Optional[float] = None,
unfreeze_end: Optional[str] = None,
unfreeze_start: Optional[str] = None,
train_bn: bool = True):
"""Unfreezes a module and adds its parameters to an optimizer."""
if (unfreeze_start is not None) or (unfreeze_end is not None):
unfreeze_modules = []
unfreeze_flag = True if unfreeze_start is None else False
# the reason for [1:] is because the named_modules return the full model
# as an unnamed nn.Sequential module as 1st member
# https://discuss.pytorch.org/t/module-children-vs-module-modules/4551/4
for name, _module in named_child_modules(module):
if unfreeze_flag:
unfreeze_modules.append(_module)
if unfreeze_start is not None and name == unfreeze_start:
unfreeze_flag = True
if unfreeze_end is not None and unfreeze_end == name:
break
module = torch.nn.Sequential(*unfreeze_modules)
_make_trainable(module)
params_lr = optimizer.param_groups[0]['lr'] if lr is None else float(lr)
optimizer.add_param_group({
'params':
filter_params(module=module, train_bn=train_bn),
'lr':
params_lr
})
def unfreeze_and_add_param_group(
modules: Union[Module, Iterable[Union[Module, Iterable]]],
optimizer: Optimizer,
lr: Optional[float] = None,
initial_denom_lr: float = 10.0,
train_bn: bool = True,
) -> None:
"""Unfreezes a module and adds its parameters to an optimizer.
Args:
modules: A module or iterable of modules to unfreeze.
Their parameters will be added to an optimizer as a new param group.
optimizer: The provided optimizer will receive new parameters and will add them to
`add_param_group`
lr: Learning rate for the new param group.
initial_denom_lr: If no lr is provided, the learning from the first param group will be used
and divided by `initial_denom_lr`.
train_bn: Whether to train the BatchNormalization layers.
"""
BaseFinetuning.make_trainable(modules)
params_lr = optimizer.param_groups[0]["lr"] if lr is None else float(
lr)
denom_lr = initial_denom_lr if lr is None else 1.0
params = BaseFinetuning.filter_params(modules,
train_bn=train_bn,
requires_grad=True)
params = BaseFinetuning.filter_on_optimizer(optimizer, params)
if params:
optimizer.add_param_group({
"params": params,
"lr": params_lr / denom_lr
})
def _unfreeze_and_add_param_group(module: Module,
optimizer: Optimizer,
lr: Optional[float] = None,
train_bn: bool = True):
"""Unfreezes a module and adds its parameters to an optimizer."""
_make_trainable(module)
params_lr = optimizer.param_groups[0]['lr'] if lr is None else float(lr)
optimizer.add_param_group(
{'params': filter_params(module=module, train_bn=train_bn),
'lr': params_lr / 10.,
})
def add_model_to_optimizer(optimizer: Optimizer, model: Module):
optimizer.add_param_group(model.parameters())
