class_docstr: Knowledge distillation loops for training and evaluation. 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
  ----------
  student_model
      The student model in knowledge distillation.
  teacher_model
      The teacher model in knowledge distillation.
  matches
      Teacher/student layer matches to compute the intermediate loss.
  critics
      The critics used in computing mutual information loss.
  baseline_funcs
      The baseline functions used in computing mutual information loss.
  hard_label_weight
      Weight for hard label loss.
  soft_label_weight
      Weight for soft label loss.
  softmax_regression_weight_label_weight
      Weight for softmax regression loss. Ref: https://www.adrianbulat.com/downloads/ICLR2021/knowledge_distillation_via_softmax_regression_representation_learning.pdf
  temperature
      A scalar to scale teacher and student logits in soft label loss.
  output_feature_loss_weight
      Weight for output_feature layer's loss.
  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.
  hard_label_loss_func
      A Pytorch loss module, e.g., nn.CrossEntropyLoss(), for hard labels.
  soft_label_loss_func
      A Pytorch loss module, e.g., nn.CrossEntropyLoss(), for soft labels.
  softmax_regression_loss_func
      A Pytorch loss module, e.g., nn.CrossEntropyLoss(), for softmax regression.
      Refer to: https://www.adrianbulat.com/downloads/ICLR2021/knowledge_distillation_via_softmax_regression_representation_learning.pdf
  output_feature_adaptor
      A Pytorch Module, e.g. nn.Linear, for adapting student output feature to the shape of teacher's.
  output_feature_loss_func
      A Pytorch loss module, e.g., nn.MSELoss(), for output_feature distance between teacher and student.
  rkd_loss_func
      A Pytorch loss module, i.e., RKDLoss in .rkd_loss, for rkd loss of output_feature between teacher and student.
  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().
  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","--","--","--","--","--","student_model","The","student","model","in","knowledge","distillation",".","teacher_model","The","teacher","model","in","knowledge","distillation",".","matches","Teacher\/student","layer","matches","to","compute","the","intermediate","loss",".","critics","The","critics","used","in","computing","mutual","information","loss",".","baseline_funcs","The","baseline","functions","used","in","computing","mutual","information","loss",".","hard_label_weight","Weight","for","hard","label","loss",".","soft_label_weight","Weight","for","soft","label","loss",".","softmax_regression_weight_label_weight","Weight","for","softmax","regression","loss",".","Ref",":","https",":","\/\/www.adrianbulat.com\/downloads\/ICLR2021\/knowledge_distillation_via_softmax_regression_representation_learning.pdf","temperature","A","scalar","to","scale","teacher","and","student","logits","in","soft","label","loss",".","output_feature_loss_weight","Weight","for","output_feature","layer","'s","loss",".","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",".","hard_label_loss_func","A","Pytorch","loss","module",",","e.g.",",","nn.CrossEntropyLoss","(",")",",","for","hard","labels",".","soft_label_loss_func","A","Pytorch","loss","module",",","e.g.",",","nn.CrossEntropyLoss","(",")",",","for","soft","labels",".","softmax_regression_loss_func","A","Pytorch","loss","module",",","e.g.",",","nn.CrossEntropyLoss","(",")",",","for","softmax","regression",".","Refer","to",":","https",":","\/\/www.adrianbulat.com\/downloads\/ICLR2021\/knowledge_distillation_via_softmax_regression_representation_learning.pdf","output_feature_adaptor","A","Pytorch","Module",",","e.g",".","nn.Linear",",","for","adapting","student","output","feature","to","the","shape","of","teacher","'s",".","output_feature_loss_func","A","Pytorch","loss","module",",","e.g.",",","nn.MSELoss","(",")",",","for","output_feature","distance","between","teacher","and","student",".","rkd_loss_func","A","Pytorch","loss","module",",","i.e.",",","RKDLoss","in",".rkd_loss",",","for","rkd","loss","of","output_feature","between","teacher","and","student",".","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","(",")",".","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,
      student_model: nn.Module,
      teacher_model: nn.Module,
      matches: List[DictConfig],
      critics: nn.ModuleList,
      baseline_funcs: nn.ModuleList,
      hard_label_weight: float,
      soft_label_weight: float,
      softmax_regression_weight: float,
      temperature: float,
      output_feature_loss_weight: float,
      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,
      hard_label_loss_func: Optional[_Loss] = None,
      soft_label_loss_func: Optional[_Loss] = None,
      softmax_regression_loss_func: Optional[_Loss] = None,
      output_feature_adaptor: Optional[nn.Module] = None,
      output_feature_loss_func: Optional[_Loss] = None,
      rkd_loss_func: Optional[nn.Module] = None,
      validation_metric: Optional[torchmetrics.Metric] = None,
      validation_metric_name: Optional[str] = None,
      custom_metric_func: Callable = None,
      test_metric: Optional[torchmetrics.Metric] = None,
      track_grad_norm: Optional[Union[int, str]] = -1,
      **kwargs,
  ):
      """
      Parameters
      ----------
      student_model
          The student model in knowledge distillation.
      teacher_model
          The teacher model in knowledge distillation.
      matches
          Teacher/student layer matches to compute the intermediate loss.
      critics
          The critics used in computing mutual information loss.
      baseline_funcs
          The baseline functions used in computing mutual information loss.
      hard_label_weight
          Weight for hard label loss.
      soft_label_weight
          Weight for soft label loss.
      softmax_regression_weight_label_weight
          Weight for softmax regression loss. Ref: https://www.adrianbulat.com/downloads/ICLR2021/knowledge_distillation_via_softmax_regression_representation_learning.pdf
      temperature
          A scalar to scale teacher and student logits in soft label loss.
      output_feature_loss_weight
          Weight for output_feature layer's loss.
      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.
      hard_label_loss_func
          A Pytorch loss module, e.g., nn.CrossEntropyLoss(), for hard labels.
      soft_label_loss_func
          A Pytorch loss module, e.g., nn.CrossEntropyLoss(), for soft labels.
      softmax_regression_loss_func
          A Pytorch loss module, e.g., nn.CrossEntropyLoss(), for softmax regression.
          Refer to: https://www.adrianbulat.com/downloads/ICLR2021/knowledge_distillation_via_softmax_regression_representation_learning.pdf
      output_feature_adaptor
          A Pytorch Module, e.g. nn.Linear, for adapting student output feature to the shape of teacher's.
      output_feature_loss_func
          A Pytorch loss module, e.g., nn.MSELoss(), for output_feature distance between teacher and student.
      rkd_loss_func
          A Pytorch loss module, i.e., RKDLoss in .rkd_loss, for rkd loss of output_feature between teacher and student.
      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().
      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.optim_type = optim_type
      self.save_hyperparameters(
          ignore=[
              "student_model",
              "teacher_model",
              "validation_metric",
              "hard_label_loss_func",
              "soft_label_loss_func",
              "custom_metric_func",
              "test_metric",
              "matches",
              "critics",
              "baseline_funcs",
              "output_feature_adaptor",
              "output_feature_loss_func",
              "rkd_loss_func",
              "loss_func",
              "mixup_fn",
          ]
      )
      if matches:
          assert len(matches) == len(critics)
          assert len(matches) == len(baseline_funcs)
      self.student_model = student_model
      self.teacher_model = teacher_model
      self.matches = matches
      self.critics = critics
      self.baseline_funcs = baseline_funcs
      self.validation_metric = validation_metric
      self.validation_metric_name = f"val_{validation_metric_name}"
      self.temperature = temperature
      self.hard_label_weight = hard_label_weight
      self.soft_label_weight = soft_label_weight
      self.softmax_regression_weight = softmax_regression_weight
      self.output_feature_loss_weight = output_feature_loss_weight
      self.hard_label_loss_func = hard_label_loss_func
      self.soft_label_loss_func = soft_label_loss_func
      self.softmax_regression_loss_func = softmax_regression_loss_func
      self.output_feature_loss_func = output_feature_loss_func
      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.output_feature_adaptor = output_feature_adaptor
      self.rkd_loss_func = rkd_loss_func
      self.track_grad_norm = track_grad_norm
  
imports: import logging
  from typing import Callable, List, Optional, Union
  import lightning.pytorch
  import torch
  import torch.nn.functional
  import torchmetrics
  from lightning.pytorch.utilities import grad_norm
  from omegaconf import DictConfig
  from torch import nn
  from torch.nn.modules.loss import _Loss
  from torchmetrics.aggregation import BaseAggregator
  from ..constants import FEATURES, LOGITS, WEIGHT
  from ..models.utils import run_model
  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 DistillerLitModule object from inherited base classes: pl with example: obj = DistillerLitModule(student_model, teacher_model, matches, critics, baseline_funcs, hard_label_weight, soft_label_weight, softmax_regression_weight, temperature, output_feature_loss_weight, optim_type, lr_choice, lr_schedule, lr, lr_decay, end_lr, lr_mult, weight_decay, warmup_steps, hard_label_loss_func, soft_label_loss_func, softmax_regression_loss_func, output_feature_adaptor, output_feature_loss_func, rkd_loss_func, validation_metric, validation_metric_name, custom_metric_func, test_metric, track_grad_norm)