mlx.nn.Module.unfreeze

mlx.nn.Module.unfreeze#

Module.unfreeze(*, recurse: bool = True, keys: str | List[str] | None = None, strict: bool = False) → Module#

Unfreeze the Module’s parameters or some of them.

This function is idempotent ie unfreezing a model that is not frozen is a noop.

Example

For instance to only train the biases of a Transformer one can do:

model = nn.Transformer()
model.freeze()
model.unfreeze(keys="bias")
Parameters:

  • recurse (bool, optional) – If True then unfreeze the parameters of the submodules as well. Default: True.

  • keys (str or list[str], optional) – If provided then only these parameters will be unfrozen otherwise all the parameters of a module. For instance unfreeze all biases by calling module.unfreeze(keys="bias").

  • strict (bool, optional) – If set to True validate that the passed keys exist. Default: False.

Returns:

The module instance after unfreezing the parameters.