class_docstr: Control the loops for training, evaluation, and prediction. This module is independent of
  the model definition. This class inherits from the Pytorch Lightning's LightningModule:
  https://lightning.ai/docs/pytorch/stable/common/lightning_module.html
docstr: Parameters
  ----------
  model
      A Pytorch model
  optim_type
      Optimizer type. We now support:
      - adamw
      - adam
      - sgd
  lr_choice
      How to set each layer's learning rate. If not specified, the default is a single
      learnng rate for all layers. Otherwise, we now support two choices:
      - two_stages
          The layers in the pretrained models have a small learning rate (lr * lr_mult),
          while the newly added head layers use the provided learning rate.
      - layerwise_decay
          The layers have decreasing learning rate from the output end to the input end.
          The intuition is that later layers are more task-related, hence larger learning rates.
  lr_schedule
      Learning rate schedule. We now support:
      - cosine_decay
          Linear warmup followed by cosine decay
      - polynomial_decay
          Linear warmup followed by polynomial decay
  lr
      Learning rate.
  lr_decay
      The learning rate decay factor (0, 1). It is used only when lr_choice is "layerwise_decay".
  end_lr
      The final learning rate after decay.
  lr_mult
      The learning rate multiplier (0, 1). It is used only when lr_choice is "two_stages".
  weight_decay
      The weight decay to regularize layer weights' l2 norm.
  warmup_steps
      How many steps to warmup learning rate. If a float (0, 1), it would represent the
      percentage of steps over all the training steps. The actual number is calculated as
      "int(warmup_steps * max_steps)". If an integer, it would be the exact step number.
  loss_func
      A Pytorch loss module, e.g., nn.CrossEntropyLoss().
  validation_metric
      A torchmetrics module used in the validation stage, e.g., torchmetrics.Accuracy().
  validation_metric_name
      Name of validation metric in case that validation_metric is a aggregation metric,
      e.g., torchmetrics.MeanMetric, whose name can't reflect the real metric name.
  custom_metric_func
      A customized metric function in case that torchmetrics doesn't have the metric.
      It is generally used together with torchmetrics' aggregators, e.g., torchmetrics.MeanMetric.
      Refer to https://github.com/PyTorchLightning/metrics/blob/master/torchmetrics/aggregation.py
  test_metric
      A torchmetrics module used in the test stage, e.g., torchmetrics.Accuracy().
  efficient_finetune
      Whether to use efficient finetuning strategies. This will be helpful for fast finetuning of large backbones.
      We support options such as:
  
      - bit_fit (only finetune the bias terms)
      - norm_fit (only finetune the weights in norm layers / bias layer)
      - lora, lora_bias, lora_norm (only finetunes decomposition matrices inserted into model, in combination with either bit_fit or norm_fit)
      - ia3, ia3_bias, ia3_norm (adds vector that scales activations by learned vectors, in combination with either bit_fit or norm_fit)
      - None (do not use efficient finetuning strategies)
  track_grad_norm
      Track the p-norm of gradients during training. May be set to ‘inf’ infinity-norm.
      If using Automatic Mixed Precision (AMP), the gradients will be unscaled before logging them.
docstr_tok: ["Parameters","--","--","--","--","--","model","A","Pytorch","model","optim_type","Optimizer","type",".","We","now","support",":","-","adamw","-","adam","-","sgd","lr_choice","How","to","set","each","layer","'s","learning","rate",".","If","not","specified",",","the","default","is","a","single","learnng","rate","for","all","layers",".","Otherwise",",","we","now","support","two","choices",":","-","two_stages","The","layers","in","the","pretrained","models","have","a","small","learning","rate","(","lr","*","lr_mult",")",",","while","the","newly","added","head","layers","use","the","provided","learning","rate",".","-","layerwise_decay","The","layers","have","decreasing","learning","rate","from","the","output","end","to","the","input","end",".","The","intuition","is","that","later","layers","are","more","task-related",",","hence","larger","learning","rates",".","lr_schedule","Learning","rate","schedule",".","We","now","support",":","-","cosine_decay","Linear","warmup","followed","by","cosine","decay","-","polynomial_decay","Linear","warmup","followed","by","polynomial","decay","lr","Learning","rate",".","lr_decay","The","learning","rate","decay","factor","(","0",",","1",")",".","It","is","used","only","when","lr_choice","is","``","layerwise_decay","''",".","end_lr","The","final","learning","rate","after","decay",".","lr_mult","The","learning","rate","multiplier","(","0",",","1",")",".","It","is","used","only","when","lr_choice","is","``","two_stages","''",".","weight_decay","The","weight","decay","to","regularize","layer","weights","'","l2","norm",".","warmup_steps","How","many","steps","to","warmup","learning","rate",".","If","a","float","(","0",",","1",")",",","it","would","represent","the","percentage","of","steps","over","all","the","training","steps",".","The","actual","number","is","calculated","as","``","int","(","warmup_steps","*","max_steps",")","''",".","If","an","integer",",","it","would","be","the","exact","step","number",".","loss_func","A","Pytorch","loss","module",",","e.g.",",","nn.CrossEntropyLoss","(",")",".","validation_metric","A","torchmetrics","module","used","in","the","validation","stage",",","e.g.",",","torchmetrics.Accuracy","(",")",".","validation_metric_name","Name","of","validation","metric","in","case","that","validation_metric","is","a","aggregation","metric",",","e.g.",",","torchmetrics.MeanMetric",",","whose","name","ca","n't","reflect","the","real","metric","name",".","custom_metric_func","A","customized","metric","function","in","case","that","torchmetrics","does","n't","have","the","metric",".","It","is","generally","used","together","with","torchmetrics","'","aggregators",",","e.g.",",","torchmetrics.MeanMetric",".","Refer","to","https",":","\/\/github.com\/PyTorchLightning\/metrics\/blob\/master\/torchmetrics\/aggregation.py","test_metric","A","torchmetrics","module","used","in","the","test","stage",",","e.g.",",","torchmetrics.Accuracy","(",")",".","efficient_finetune","Whether","to","use","efficient","finetuning","strategies",".","This","will","be","helpful","for","fast","finetuning","of","large","backbones",".","We","support","options","such","as",":","-","bit_fit","(","only","finetune","the","bias","terms",")","-","norm_fit","(","only","finetune","the","weights","in","norm","layers","\/","bias","layer",")","-","lora",",","lora_bias",",","lora_norm","(","only","finetunes","decomposition","matrices","inserted","into","model",",","in","combination","with","either","bit_fit","or","norm_fit",")","-","ia3",",","ia3_bias",",","ia3_norm","(","adds","vector","that","scales","activations","by","learned","vectors",",","in","combination","with","either","bit_fit","or","norm_fit",")","-","None","(","do","not","use","efficient","finetuning","strategies",")","track_grad_norm","Track","the","p-norm","of","gradients","during","training",".","May","be","set","to","\u2018","inf","\u2019","infinity-norm",".","If","using","Automatic","Mixed","Precision","(","AMP",")",",","the","gradients","will","be","unscaled","before","logging","them","."]
code: def __init__(
      self,
      model: nn.Module,
      optim_type: Optional[str] = None,
      lr_choice: Optional[str] = None,
      lr_schedule: Optional[str] = None,
      lr: Optional[float] = None,
      lr_decay: Optional[float] = None,
      end_lr: Optional[Union[float, int]] = None,
      lr_mult: Optional[Union[float, int]] = None,
      weight_decay: Optional[float] = None,
      warmup_steps: Optional[int] = None,
      loss_func: Optional[_Loss] = None,
      validation_metric: Optional[torchmetrics.Metric] = None,
      validation_metric_name: Optional[str] = None,
      custom_metric_func: Callable = None,
      test_metric: Optional[torchmetrics.Metric] = None,
      efficient_finetune: Optional[str] = None,
      trainable_param_names: Optional[List] = None,
      mixup_fn: Optional[MixupModule] = None,
      mixup_off_epoch: Optional[int] = 0,
      model_postprocess_fn: Callable = None,
      skip_final_val: Optional[bool] = False,
      track_grad_norm: Optional[Union[int, str]] = -1,
  ):
      """
      Parameters
      ----------
      model
          A Pytorch model
      optim_type
          Optimizer type. We now support:
          - adamw
          - adam
          - sgd
      lr_choice
          How to set each layer's learning rate. If not specified, the default is a single
          learnng rate for all layers. Otherwise, we now support two choices:
          - two_stages
              The layers in the pretrained models have a small learning rate (lr * lr_mult),
              while the newly added head layers use the provided learning rate.
          - layerwise_decay
              The layers have decreasing learning rate from the output end to the input end.
              The intuition is that later layers are more task-related, hence larger learning rates.
      lr_schedule
          Learning rate schedule. We now support:
          - cosine_decay
              Linear warmup followed by cosine decay
          - polynomial_decay
              Linear warmup followed by polynomial decay
      lr
          Learning rate.
      lr_decay
          The learning rate decay factor (0, 1). It is used only when lr_choice is "layerwise_decay".
      end_lr
          The final learning rate after decay.
      lr_mult
          The learning rate multiplier (0, 1). It is used only when lr_choice is "two_stages".
      weight_decay
          The weight decay to regularize layer weights' l2 norm.
      warmup_steps
          How many steps to warmup learning rate. If a float (0, 1), it would represent the
          percentage of steps over all the training steps. The actual number is calculated as
          "int(warmup_steps * max_steps)". If an integer, it would be the exact step number.
      loss_func
          A Pytorch loss module, e.g., nn.CrossEntropyLoss().
      validation_metric
          A torchmetrics module used in the validation stage, e.g., torchmetrics.Accuracy().
      validation_metric_name
          Name of validation metric in case that validation_metric is a aggregation metric,
          e.g., torchmetrics.MeanMetric, whose name can't reflect the real metric name.
      custom_metric_func
          A customized metric function in case that torchmetrics doesn't have the metric.
          It is generally used together with torchmetrics' aggregators, e.g., torchmetrics.MeanMetric.
          Refer to https://github.com/PyTorchLightning/metrics/blob/master/torchmetrics/aggregation.py
      test_metric
          A torchmetrics module used in the test stage, e.g., torchmetrics.Accuracy().
      efficient_finetune
          Whether to use efficient finetuning strategies. This will be helpful for fast finetuning of large backbones.
          We support options such as:
  
          - bit_fit (only finetune the bias terms)
          - norm_fit (only finetune the weights in norm layers / bias layer)
          - lora, lora_bias, lora_norm (only finetunes decomposition matrices inserted into model, in combination with either bit_fit or norm_fit)
          - ia3, ia3_bias, ia3_norm (adds vector that scales activations by learned vectors, in combination with either bit_fit or norm_fit)
          - None (do not use efficient finetuning strategies)
      track_grad_norm
          Track the p-norm of gradients during training. May be set to ‘inf’ infinity-norm.
          If using Automatic Mixed Precision (AMP), the gradients will be unscaled before logging them.
  
      """
      super().__init__()
      self.save_hyperparameters(
          ignore=[
              "model",
              "validation_metric",
              "test_metric",
              "loss_func",
              "model_postprocess_fn",
              "mixup_fn",
              "trainable_param_names",
          ]
      )
      self.model = model
      self.validation_metric = validation_metric
      self.validation_metric_name = f"val_{validation_metric_name}"
      self.loss_func = loss_func
      self.mixup_fn = mixup_fn
      if (
          isinstance(validation_metric, BaseAggregator)
          and custom_metric_func is None
      ):
          raise ValueError(
              f"validation_metric {validation_metric} is an aggregation metric,"
              "which must be used with a customized metric function."
          )
      self.custom_metric_func = custom_metric_func
      self.model_postprocess_fn = model_postprocess_fn
      self.trainable_param_names = (
          trainable_param_names if trainable_param_names else []
      )
      self.skip_final_val = skip_final_val
      self.track_grad_norm = track_grad_norm
  
imports: import logging
  from typing import Callable, Dict, List, Optional, Union
  import lightning.pytorch
  import torch
  import torch.nn.functional
  import torchmetrics
  from lightning.pytorch.strategies import DeepSpeedStrategy
  from lightning.pytorch.utilities import grad_norm
  from torch import nn
  from torch.nn.modules.loss import _Loss
  from torchmetrics.aggregation import BaseAggregator
  from ..constants import LM_TARGET, LOGITS, T_FEW, TEMPLATE_LOGITS, WEIGHT
  from ..data.mixup import MixupModule, multimodel_mixup
  from ..models.utils import run_model
  from .semantic_seg_metrics import COD, Balanced_Error_Rate
  from .utils import apply_layerwise_lr_decay, apply_single_lr, apply_two_stages_lr, get_lr_scheduler, get_optimizer
label_desc: Use image node_id 1 to create a new LitModule object from inherited base classes: pl with example: obj = LitModule(model, optim_type, lr_choice, lr_schedule, lr, lr_decay, end_lr, lr_mult, weight_decay, warmup_steps, loss_func, validation_metric, validation_metric_name, custom_metric_func, test_metric, efficient_finetune, trainable_param_names, mixup_fn, mixup_off_epoch, model_postprocess_fn, skip_final_val, track_grad_norm)