frarch.utils.freezable_module_dict module#

class frarch.utils.freezable_module_dict.FreezableModuleDict(freeze: Optional[List[str]] = None, **kwargs)[source]#

Bases: torch.nn.modules.container.ModuleDict

__init__(freeze: Optional[List[str]] = None, **kwargs)[source]#

Torch ModuleDict with freeze capabilities.

Parameters

  • freeze (Optional[List[str]]) – list of names of the parameter’s to freeze.

  • unfreeze (Optional[List[str]]) – list of names of the parameter’s to unfreeze.

  • **kwargs – ModuleDict arguments.

Raises

FreezeError – both freeze and unfreeze can’t be set simultaneously.