class_docstr: Model and training loop for the DVC model.
  
  Combines a pre-defined DVC model with its training loop for use with
  PyTorch Lightning.
  
  Args:
      model: The DVC model to train.
      training_stage: Current stage of training process. One of
          ``("1_motion_estimation", "2_motion_compression",
          "3_motion_compensation", "4_total")``. See DVC paper for details.
      pretrained_model_name: Name of model from CompressAI model zoo to use
          for compressing I-frames.
      pretrained_model_quality_level: Quality level of model from CompressAI.
      num_pframes: Number of P-frames to process for training.
      distortion_type: Type of distortion loss function. Must be from
          ``("MSE")``.
      distortion_lambda: A scaling factor for the distortion term of the
          loss.
      learning_rate: passed to the main network optimizer (i.e. the one that
          adjusts the analysis and synthesis parameters).
      aux_learning_rate: passed to the optimizer that learns the quantiles
          used to build the CDF table for the entropy codder.
      lr_scheduler_params: Used for ``StepLR``, specify ``step_size`` and
          ``gamma``.
class_docstr_tok: ["Model","and","training","loop","for","the","DVC","model",".","Combines","a","pre-defined","DVC","model","with","its","training","loop","for","use","with","PyTorch","Lightning",".","Args",":","model",":","The","DVC","model","to","train",".","training_stage",":","Current","stage","of","training","process",".","One","of","``","(","``","1_motion_estimation","''",",","``","2_motion_compression","''",",","``","3_motion_compensation","''",",","``","4_total","''",")","``",".","See","DVC","paper","for","details",".","pretrained_model_name",":","Name","of","model","from","CompressAI","model","zoo","to","use","for","compressing","I-frames",".","pretrained_model_quality_level",":","Quality","level","of","model","from","CompressAI",".","num_pframes",":","Number","of","P-frames","to","process","for","training",".","distortion_type",":","Type","of","distortion","loss","function",".","Must","be","from","``","(","``","MSE","''",")","``",".","distortion_lambda",":","A","scaling","factor","for","the","distortion","term","of","the","loss",".","learning_rate",":","passed","to","the","main","network","optimizer","(","i.e",".","the","one","that","adjusts","the","analysis","and","synthesis","parameters",")",".","aux_learning_rate",":","passed","to","the","optimizer","that","learns","the","quantiles","used","to","build","the","CDF","table","for","the","entropy","codder",".","lr_scheduler_params",":","Used","for","``","StepLR","``",",","specify","``","step_size","``","and","``","gamma","``","."]
code: def __init__(
      self,
      model: DVC,
      training_stage: str,
      pretrained_model_name: str,
      pretrained_model_quality_level: int,
      num_pframes: int = 1,
      distortion_type: str = "MSE",
      distortion_lambda: float = 256.0,
      learning_rate: float = 1e-4,
      aux_learning_rate: float = 1e-3,
      lr_scheduler_params: Optional[Tuple[int, float]] = None,
      grad_clip_value: float = 1.0,
  ):
      super().__init__()
      self.automatic_optimization = False
      self.training_stage = training_stage
      self.learning_rate = learning_rate
      self.num_pframes = num_pframes
      self.aux_learning_rate = aux_learning_rate
      self.distortion_lambda = distortion_lambda
      self.lr_scheduler_params = lr_scheduler_params
      self.grad_clip_value = grad_clip_value
      self.iframe_model = models[pretrained_model_name](
          pretrained_model_quality_level, pretrained=True
      )
  
      # make sure training stage is valid
      if training_stage not in TRAINING_STAGES.keys():
          raise ValueError(
              f"training stage {training_stage} not recognized."
          )
      if distortion_type == "MSE":
          distortion_loss = torch.nn.MSELoss()
      else:
          raise ValueError(
              f"distortion_type {distortion_type} not recognized."
          )
  
      # set up loss functions
      def distortion_fn(image1, image2):
          return self.distortion_lambda * distortion_loss(
              image1, image2
          )
  
      def entropy_fn(probabilities, reduction_factor):
          return (
              ncF.information_content(probabilities, reduction="sum")
              / reduction_factor
          )
  
      self.loss_functions = LossFunctions(distortion_fn, entropy_fn)
  
      # set up model - this includes functions for computing losses and
      # compressing
      self.model: DvcStage = TRAINING_STAGES[training_stage](
          model, num_pframes, self.loss_functions
      )
  
      # metrics
      self.psnr = torchmetrics.PSNR(data_range=1.0)
  
imports: from typing import NamedTuple, Optional, Tuple
  import torch
  import torch.optim
  import torchmetrics
  from _utils import DvcStage, DvcStage1, DvcStage2, DvcStage3, DvcStage4and5, LossFunctions
  from compressai.zoo import models
  from pytorch_lightning import LightningModule
  from torch import Tensor
  import neuralcompression.functional
  from neuralcompression.models import DVC
label_desc: Use image node_id 1 to create a new DvcModule object from inherited base classes: LightningModule with example: obj = DvcModule(model, training_stage, pretrained_model_name, pretrained_model_quality_level, num_pframes, distortion_type, distortion_lambda, learning_rate, aux_learning_rate, lr_scheduler_params, grad_clip_value)