class_docstr: Standard solver for AudioCraft.
  
  The standard solver implements a base training loop with the following stages:
  train, valid, evaluate and generate that are expected to be all defined for
  solvers in AudioCraft. It also provides a nice default management of Dora history replay,
  checkpoint management across epoch, and logging configuration.
  
  AudioCraft solvers must inherit from the StandardSolver and define the methods
  associated to each stage as well as the show, build_model and build_dataloaders methods.
class_docstr_tok: ["Standard","solver","for","AudioCraft",".","The","standard","solver","implements","a","base","training","loop","with","the","following","stages",":","train",",","valid",",","evaluate","and","generate","that","are","expected","to","be","all","defined","for","solvers","in","AudioCraft",".","It","also","provides","a","nice","default","management","of","Dora","history","replay",",","checkpoint","management","across","epoch",",","and","logging","configuration",".","AudioCraft","solvers","must","inherit","from","the","StandardSolver","and","define","the","methods","associated","to","each","stage","as","well","as","the","show",",","build_model","and","build_dataloaders","methods","."]
code: def __init__(self, cfg: omegaconf.DictConfig):
      super().__init__()
      self.logger.info(
          f"Instantiating solver {self.__class__.__name__} for XP {self.xp.sig}"
      )
      self.logger.info(f"All XP logs are stored in {self.xp.folder}")
      self.cfg = cfg
      self.device = cfg.device
      self.model: nn.Module
      self._continue_best_source_keys = [
          "best_state",
          "fsdp_best_state",
      ]
      self._fsdp_modules: tp.List[fsdp.FSDP] = []
      self._ema_sources: nn.ModuleDict = nn.ModuleDict()
      self.ema: tp.Optional[optim.ModuleDictEMA] = None
      self.dataloaders: tp.Dict[
          str, torch.utils.data.DataLoader
      ] = dict()
      self._log_updates = self.cfg.logging.get("log_updates", 10)
      if self.cfg.logging.log_tensorboard:
          self.init_tensorboard(**self.cfg.get("tensorboard"))
      if self.cfg.logging.log_wandb and self:
          self.init_wandb(**self.cfg.get("wandb"))
      # keep a copy of the best performing state for stateful objects
      # used for evaluation and generation stages
      dtype_best: tp.Optional[torch.dtype] = None
      if self.cfg.fsdp.use:
          dtype_best = getattr(torch, self.cfg.fsdp.param_dtype)  # type: ignore
          assert isinstance(dtype_best, torch.dtype)
      elif self.cfg.autocast:
          dtype_best = getattr(torch, self.cfg.autocast_dtype)  # type: ignore
          assert isinstance(dtype_best, torch.dtype)
      self.best_state: BestStateDictManager = BestStateDictManager(
          dtype=dtype_best
      )
      # Hacky support for keeping a copy of the full best state in rank0.
      self.fsdp_best_state: tp.Dict[str, tp.Any] = {}
      self.register_stateful(
          "best_state", "fsdp_best_state"
      )  # register best_state object to keep it in state_dict
      self._new_best_state: bool = False  # should save a new checkpoint
      # instantiate datasets and appropriate number of updates per epoch
      self.build_dataloaders()
      if self.cfg.execute_only is None:
          assert (
              "train" in self.dataloaders
          ), "The train dataset split must be provided."
          assert (
              "valid" in self.dataloaders
          ), "The valid dataset split must be provided."
      self.train_updates_per_epoch = (
          len(self.dataloaders["train"])
          if "train" in self.dataloaders
          else 0
      )
      if self.cfg.optim.updates_per_epoch:
          self.train_updates_per_epoch = (
              self.cfg.optim.updates_per_epoch
          )
      self.total_updates = (
          self.train_updates_per_epoch * self.cfg.optim.epochs
      )
      # instantiate model & exponential moving average on the model
      self.build_model()
      self.logger.info("Model hash: %s", model_hash(self.model))
      assert (
          "model" in self.stateful.sources
      ), "Please register the model to stateful with self.register_stateful('model') in build_model."
      self.profiler = Profiler(self.model, **self.cfg.profiler)
      self.initialize_ema()
      self.register_stateful("ema")
      assert (
          self.ema is None or "ema" in self.stateful.sources
      ), "Please register the ema to stateful with self.register_stateful('ema') in build_model."
      self.deadlock_detect = DeadlockDetect(**self.cfg.deadlock)
      # basic statistics on the trained model
      model_size = (
          sum(
              p.numel()
              for p in self.model.parameters()
              if p.requires_grad
          )
          / 1e6
      )
      # one copy of grad, one copy of momentum, one copy of denominator and model weights.
      # and 4 bytes for each float!
      mem_usage = model_size * 4 * 4 / 1000
      self.logger.info("Model size: %.2f M params", model_size)
      self.logger.info(
          "Base memory usage, with model, grad and optim: %.2f GB",
          mem_usage,
      )
  
imports: from abc import ABC, abstractmethod
  from contextlib import contextmanager
  from pathlib import Path
  import typing
  import flashy
  import omegaconf
  import torch
  from torch import nn
  from ..None import optim
  from ..optim import fsdp
  from ..utils import checkpoint
  from ..utils.autocast import TorchAutocast
  from ..utils.best_state import BestStateDictManager
  from ..utils.deadlock import DeadlockDetect
  from ..utils.profiler import Profiler
  from ..utils.utils import copy_state, dict_from_config, model_hash, with_rank_rng
label_desc: Use image node_id 1 to create a new StandardSolver object from inherited base classes: ABC, flashy with example: obj = StandardSolver(cfg)