code: def build_optimizer_param_groups(
      model: nn.Module,
      learning_rate: Union[float, Mapping[str, float]],
      weight_decay: Union[float, Mapping[str, float]],
      apply_weight_decay_on_bias: bool = True,
      apply_weight_decay_on_norm: bool = True,
  ) -> Tuple[List[Dict], Dict]:
      """
  
      Args:
          model: A model whose parameters will be used to fill corresponding parameter groups.
          learning_rate: A single number of dictionary of layer-wise learning rate parameters.
          weight_decay: A single number of dictionary of layer-wise weight decay parameters.
          apply_weight_decay_on_bias: If True, weight decay is applied to bias on Linear & Conv layers (default).
              This parameter is False, it overrules the matching layer-wise weight-decay parameter.
          apply_weight_decay_on_norm: If True, weight decay is applied normalization layers (default).
              This parameter is False, it overrules the matching layer-wise weight-decay parameter.
  
      Returns:
  
      """
      model = get_non_wrapped_model(model)
  
      if (
          isinstance(learning_rate, Mapping)
          and "_default_" not in learning_rate
      ):
          raise RuntimeError(
              "When using layerwise learning rate, a key _default_ must be present to indicate default LR"
          )
  
      if (
          isinstance(weight_decay, Mapping)
          and "_default_" not in weight_decay
      ):
          raise RuntimeError(
              "When using layerwise weight decay, a key _default_ must be present to indicate default LR"
          )
  
      if isinstance(learning_rate, numbers.Number):
          learning_rate = {"_default_": float(learning_rate)}
      if isinstance(weight_decay, numbers.Number):
          weight_decay = {"_default_": float(weight_decay)}
  
      default_learning_rate = float(learning_rate["_default_"])
      default_weight_decay = float(weight_decay["_default_"])
  
      learning_rate = list(learning_rate.items())
      weight_decay = list(weight_decay.items())
  
      parameter_groups = {}
  
      norm_layers = (
          nn.BatchNorm1d,
          nn.BatchNorm2d,
          nn.BatchNorm3d,
          nn.GroupNorm,
          nn.LayerNorm,
          nn.InstanceNorm1d,
          nn.InstanceNorm3d,
          nn.InstanceNorm2d,
          nn.SyncBatchNorm,
      )
  
      layers_with_bias = (
          nn.Linear,
          nn.Conv1d,
          nn.Conv2d,
          nn.Conv3d,
          nn.ConvTranspose1d,
          nn.ConvTranspose2d,
          nn.ConvTranspose3d,
      )
  
      def get_param_group(parameter_name: str, module: nn.Module):
          matching_lr_index = "default"
          matching_lr_value = default_learning_rate
          for lr_index, (prefix, lr) in enumerate(learning_rate):
              if parameter_name.startswith(prefix):
                  matching_lr_index = prefix
                  matching_lr_value = lr
                  break
  
          matching_wd_index = "default"
          matching_wd_value = default_weight_decay
          for wd_index, (prefix, wd) in enumerate(weight_decay):
              if parameter_name.startswith(prefix):
                  matching_wd_index = prefix
                  matching_wd_value = wd
                  break
  
          if not apply_weight_decay_on_norm and isinstance(
              module, norm_layers
          ):
              matching_wd_value = 0
              matching_wd_index = "no_wd_on_norm"
  
          if (
              not apply_weight_decay_on_bias
              and isinstance(module, layers_with_bias)
              and parameter_name.endswith(".bias")
          ):
              matching_wd_value = 0
              matching_wd_index = "no_wd_on_bias"
  
          if matching_lr_index == matching_wd_index:
              param_group_name = f"{matching_lr_index}"
          else:
              param_group_name = (
                  f"{matching_lr_index}_{matching_wd_index}"
              )
          if param_group_name not in parameter_groups:
              parameter_groups[param_group_name] = ParametersGroup(
                  lr=matching_lr_value,
                  weight_decay=matching_wd_value,
                  name=param_group_name,
                  params=[],
              )
  
          return parameter_groups[param_group_name]
  
      # All optimizable parameters
      parameters = get_named_optimizable_parameters(model)
      total_optimizable_params = 0
  
      for parameter_name, parameter in parameters:
          total_optimizable_params += parameter.numel()
          module_name = ".".join(parameter_name.split(".")[:-1])
          module = recursive_getattr(model, module_name)
  
          param_group: ParametersGroup = get_param_group(
              parameter_name, module
          )
          param_group.params.append(parameter)
  
      defaults = {
          "lr": default_learning_rate,
          "weight_decay": default_weight_decay,
      }
      param_groups = [x.asdict() for x in parameter_groups.values()]
  
      total_params_count_from_groups = 0
      for pg in param_groups:
          total_params_count_from_groups += sum(
              x.numel() for x in pg["params"]
          )
  
      if total_params_count_from_groups != total_optimizable_params:
          raise RuntimeError(
              f"Detected mismatch in total number of optimizable parameters ({total_optimizable_params}) and"
              f"number of parameters across each groups ({total_params_count_from_groups})."
              f"This is likely indicate a bug in build_optimizer_param_groups."
              f"Please report a bug to https://github.com/BloodAxe/pytorch-toolbelt/issues/new?assignees=&labels=&template=bug-report.md"
          )
      return param_groups, defaults
  
file_path: pytorch-toolbelt/pytorch_toolbelt/optimization/functional.py
imports: import dataclasses
  import numbers
  from typing import Optional, Iterator, Dict, Union, List, Tuple, Mapping
  from torch import nn
  from pytorch_toolbelt.utils.distributed import get_rank, get_world_size, is_dist_avail_and_initialized
  from pytorch_toolbelt.utils.torch_utils import get_non_wrapped_model
label_desc: Use image node_id 3 for calling a global function with example usage: build_optimizer_param_groups(model, learning_rate, weight_decay, apply_weight_decay_on_bias, apply_weight_decay_on_norm) and returns: param_groups,  defaults, parameter_groups