aiaccel.torch.lightning.build_param_groups

aiaccel.torch.lightning.build_param_groups(named_params: Iterator[tuple[str, Parameter]], groups: list[dict[str, Any]]) → list[dict[str, Any]]

Build parameter groups for the optimizer based on the provided patterns.

Parameters:

• named_params (Iterator[tuple[str, nn.Parameter]]) – An iterator of named parameters.

• groups (list[dict[str, Any]]) – A list of dictionaries where each dictionary contains a “pattern” key that specifies the parameter names to match (fnmatch), and other optional keys.

Example: In your config file, you might have:

optimizer_config:
  _target_: aiaccel.torch.lightning.OptimizerConfig
  optimizer_generator:
    _partial_: True
    _target_: torch.optim.AdamW
    weight_decay: 0.01
  params_transformer:
      _partial_: True
      _target_: aiaccel.torch.lightning.build_param_groups
      groups:
        - pattern: "*bias"
          lr: 0.01
        - pattern: "*weight"
          lr: 0.001

This will create two parameter groups: one for biases with a learning rate of 0.01 and another for weights with a learning rate of 0.001.