xnn.common.finetune.freeze.freeze_parameters

xnn.common.finetune.freeze.freeze_parameters(model, freeze=(), train_only=())

Set requires_grad of the parameters of model by name pattern.

Parameters:

• model (torch.nn.Module) – The model.

• freeze (sequence of str, optional) – fnmatch patterns of parameter names to freeze (["model.interactions.0.*", "*node_embedding*"]).

• train_only (sequence of str, optional) – Patterns of the parameters to keep trainable; every other parameter is frozen (["*readouts*", "*atom_ref*"] trains the readout only). freeze is applied afterwards.

Returns:

(trainable, total) parameter counts after the change.

Return type:

tuple of int