nip.utils.training.ParamGroupFreezer

class nip.utils.training.ParamGroupFreezer(optimizer: Optimizer, param_group_collections: dict[str, Iterable[dict[str, Any]]], use_required_grad: bool = True)

A class for freezing groups of parameters during training.

The param_group_groups collects the optimizer’s parameter groups into named collections. Each collection can be frozen and unfrozen at any point.

Freezing is achieved by setting removing the parameter groups from the optimizer and optionally setting the requires_grad attribute to False.

Parameters:

• optimizer (Optimizer) – The optimizer which has the parameter groups to freeze.

• param_group_collections (dict[str, Iterable[dict[str, Any]]]) – A dictionary whose keys name collections of parameter groups which can be frozen during training. The values are lists of dictionaries which define the parameter groups (as stored in optimizer.param_groups). The keys of the dictionaries are “params” and “lr”.

• use_required_grad (bool, default=True) – Whether to set the requires_grad to False when freezing the parameters.

Methods Summary

__init__(optimizer, param_group_collections)
freeze(collection_name)	Freeze the parameter groups in the named collection.
unfreeze(collection_name)	Unfreeze the parameter groups in the named collection.

Methods

__init__(optimizer: Optimizer, param_group_collections: dict[str, Iterable[dict[str, Any]]], use_required_grad: bool = True)

freeze(collection_name: str)

Freeze the parameter groups in the named collection.

Parameters:

collection_name (str) – The name of the collection of parameter groups to freeze.

unfreeze(collection_name: str)

Unfreeze the parameter groups in the named collection.

Parameters:

collection_name (str) – The name of the collection of parameter groups to unfreeze.