def separate_modules_for_weight_decay(
    network,
    whitelist_weight_modules,
    blacklist_weight_modules,
    layers_to_exclude
) -> tuple[set[str], set[str]]

Separate the parameters of network into two sets: one set of parameters that will have weight decay, and one set that will not.

Arguments:

• network(torch.nn.Module): Network whose modules we want to separate.
• whitelist_weight_modules(tuple[Type[torch.nn.Module], ...]): Modules that should have weight decay applied to the weights.
• blacklist_weight_modules(tuple[Type[torch.nn.Module], ...]): Modules that should not have weight decay applied to the weights.
• layers_to_exclude(set[str] | None): Names of layers that should be excluded. Defaults to None. (default: None)

Returns:

• Sets of modules with and without weight decay.

Modifies AdamW (Adam with weight decay) to not apply weight decay on the bias and layer normalization weights, and optionally additional modules.

Methods

def __init__(
    self,
    network,
    lr,
    weight_decay,
    whitelist_weight_modules,
    blacklist_weight_modules,
    layers_to_exclude
) -> None