code: def apply_two_stages_lr(
      model: nn.Module,
      lr: float,
      lr_mult: Union[float, int],
      weight_decay: float,
      return_params: Optional[bool] = True,
  ):
      """
      Set up the pretrained backbone to use a smaller learning rate (lr * lr_mult).
      The newly added head layers use the normal learning rate (lr).
      Layer normalization parameters and other layers' bias parameters don't use weight decay.
  
      Parameters
      ----------
      model
          A Pytorch model.
      lr
          The learning rate.
      lr_mult
          The multiplier (0, 1) to scale down the learning rate.
      weight_decay
          Weight decay.
      return_params
          return_params
          Whether to return parameters or their names. If you want to double-check
          whether the learning rate setup is as expected, you can set "return_params=False",
          and print the layer names along with their learning rates through
          "print("Param groups = %s" % json.dumps(optimizer_grouped_parameters, indent=2))".
  
      Returns
      -------
      The grouped parameters or their names.
      """
      decay_param_names = get_weight_decay_param_names(model)
  
      optimizer_grouped_parameters = [
          {
              "params": [
                  p if return_params else n
                  for n, p in model.named_parameters()
                  if n in decay_param_names
                  and not any(bb in n for bb in model.head_layer_names)
              ],
              "weight_decay": weight_decay,
              "lr": lr,
          },
          {
              "params": [
                  p if return_params else n
                  for n, p in model.named_parameters()
                  if n not in decay_param_names
                  and not any(bb in n for bb in model.head_layer_names)
              ],
              "weight_decay": 0.0,
              "lr": lr,
          },
          {
              "params": [
                  p if return_params else n
                  for n, p in model.named_parameters()
                  if n in decay_param_names
                  and any(bb in n for bb in model.head_layer_names)
              ],
              "weight_decay": weight_decay,
              "lr": lr * lr_mult,
          },
          {
              "params": [
                  p if return_params else n
                  for n, p in model.named_parameters()
                  if n not in decay_param_names
                  and any(bb in n for bb in model.head_layer_names)
              ],
              "weight_decay": 0.0,
              "lr": lr * lr_mult,
          },
      ]
  
      return optimizer_grouped_parameters
  
imports: import functools
  import logging
  import re
  import warnings
  from typing import Any, Dict, List, Optional, Tuple, Union
  import torch
  import torchmetrics
  from omegaconf import DictConfig, OmegaConf
  from packaging import version
  from pytorch_metric_learning import distances, losses, miners
  from torch import nn, optim
  from torch.nn import functional
  from torchmetrics.detection.mean_ap import MeanAveragePrecision
  from transformers import Adafactor
  from transformers.models.mask2former.modeling_mask2former import Mask2FormerConfig, Mask2FormerLoss
  from transformers.trainer_pt_utils import get_parameter_names
  from ..constants import ACC, ACCURACY, AVERAGE_PRECISION, BER, BINARY, BIT_FIT, COLUMN_FEATURES, CONTRASTIVE_LOSS, COSINE_EMBEDDING_LOSS, COSINE_SIMILARITY, CROSS_ENTROPY, DETECTION_METRICS, DIRECT_LOSS, EM, F1, F1_MACRO, F1_MICRO, F1_WEIGHTED, FEATURES, FEW_SHOT_CLASSIFICATION, FM, HIT_RATE, IA3, IA3_BIAS, IA3_LORA, IA3_LORA_BIAS, IA3_LORA_NORM, IA3_NORM, IOU, LOG_LOSS, LORA, LORA_BIAS, LORA_NORM, MAE, MULTI_NEGATIVES_SOFTMAX_LOSS, MULTICLASS, NER, NER_TOKEN_F1, NORM_FIT, OBJECT_DETECTION, OPEN_VOCABULARY_OBJECT_DETECTION, OVERALL_ACCURACY, PAIR_MARGIN_MINER, PEARSONR, PEFT_STRATEGIES, QUADRATIC_KAPPA, R2, RECALL, REGRESSION, RMSE, ROC_AUC, ROOT_MEAN_SQUARED_ERROR, SEMANTIC_SEGMENTATION, SM, SPEARMANR
  from .losses import BBCEWithLogitLoss, FocalLoss, MultiNegativesSoftmaxLoss, SoftTargetCrossEntropy, StructureLoss
  from .lr_scheduler import get_cosine_schedule_with_warmup, get_linear_schedule_with_warmup, get_polynomial_decay_schedule_with_warmup
  from .semantic_seg_metrics import COD_METRICS_NAMES, Balanced_Error_Rate, Binary_IoU, Multiclass_IoU
label_desc: Use image node_id 9 for calling a global function with example usage: apply_two_stages_lr(model, lr, lr_mult, weight_decay, return_params) and returns: optimizer_grouped_parameters