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 validation_step(self, batch, batch_idx):
      log_key = f"{self.training_stage}/val_"
      # gop = "Group of Pictures"
      logging_metrics = LoggingMetrics(
          gop_total_loss=0,
          gop_distortion_loss=0,
          gop_bpp=0,
          gop_flow_bpp=0,
          gop_residual_bpp=0,
      )
      image2_list = []
      image2_est_list = []
  
      batch, gop_bpp = self.compress_iframe(
          batch
      )  # bpp_total w/o grads
      image1 = batch[:, 0]
      for i in range(self.num_pframes):
          image2 = batch[:, i + 1]
          loss_values, images = self.model.compute_batch_loss(
              image1, image2
          )
          image1 = images.image2_est  # images are detached
  
          # keep track of these for other distortion metrics
          image2_list.append(images.image2)
          image2_est_list.append(images.image2_est)
  
          # loss function collection
          loss, logging_metrics = self.compute_loss_and_metrics(
              loss_values, logging_metrics
          )
  
      # stat reductions and logging
      reduction = self.num_pframes + 1
      self.log("val_loss", loss)
      self.log_all_metrics(
          log_key,
          reduction,
          image2_list,
          image2_est_list,
          logging_metrics,
      )
  
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 9 for calling the DvcModule obj's underlying member method code with example usage: obj.validation_step(batch, batch_idx) without return types