content: from typing import Union
  
  import torch
  from torch._dynamo import OptimizedModule
  
  from nnunetv2.training.lr_scheduler.warmup import Lin_incr_LRScheduler, PolyLRScheduler_offset
  from nnunetv2.training.nnUNetTrainer.nnUNetTrainer import nnUNetTrainer
  from torch.nn.parallel import DistributedDataParallel as DDP
  
  from nnunetv2.utilities.helpers import empty_cache
  
  
  class nnUNetTrainer_warmup(nnUNetTrainer):
      """
      Does a warmup of the entire architecture
      Then does normal training
      """
  
      def __init__(
          self,
          plans: dict,
          configuration: str,
          fold: int,
          dataset_json: dict,
          device: torch.device = torch.device("cuda"),
      ):
          super().__init__(plans, configuration, fold, dataset_json, device)
          #### hyperparameters for warmup
          self.warmup_duration_whole_net = 50  # lin increase whole network
          self.num_epochs = 1000
          self.training_stage = None  # 'warmup_all', 'train'
  
      def configure_optimizers(self, stage: str = "warmup_all"):
          assert stage in ["warmup_all", "train"]
  
          if self.training_stage == stage:
              return self.optimizer, self.lr_scheduler
  
          if isinstance(self.network, DDP):
              params = self.network.module.parameters()
          else:
              params = self.network.parameters()
  
          if stage == "warmup_all":
              self.print_to_log_file("train whole net, warmup")
              optimizer = torch.optim.SGD(
                  params, self.initial_lr, weight_decay=self.weight_decay, momentum=0.99, nesterov=True
              )
              lr_scheduler = Lin_incr_LRScheduler(optimizer, self.initial_lr, self.warmup_duration_whole_net)
              self.print_to_log_file(f"Initialized warmup_all optimizer and lr_scheduler at epoch {self.current_epoch}")
          else:
              self.print_to_log_file("train whole net, default schedule")
              if self.training_stage == "warmup_all":
                  # we can keep the existing optimizer and don't need to create a new one. This will allow us to keep
                  # the accumulated momentum terms which already point in a useful driection
                  optimizer = self.optimizer
              else:
                  optimizer = torch.optim.SGD(
                      params, self.initial_lr, weight_decay=self.weight_decay, momentum=0.99, nesterov=True
                  )
              lr_scheduler = PolyLRScheduler_offset(
                  optimizer, self.initial_lr, self.num_epochs, self.warmup_duration_whole_net
              )
              self.print_to_log_file(f"Initialized train optimizer and lr_scheduler at epoch {self.current_epoch}")
          self.training_stage = stage
          empty_cache(self.device)
          return optimizer, lr_scheduler
  
      def on_train_epoch_start(self):
          if self.current_epoch == 0:
              self.optimizer, self.lr_scheduler = self.configure_optimizers("warmup_all")
          elif self.current_epoch == self.warmup_duration_whole_net:
              self.optimizer, self.lr_scheduler = self.configure_optimizers("train")
  
          super().on_train_epoch_start()
  
      def load_checkpoint(self, filename_or_checkpoint: Union[dict, str]) -> None:
          """
          We need to overwrite that entire function because we need to fiddle the correct optimizer in between
          loading the checkpoint and applying the optimizer states. Yuck.
          """
          if not self.was_initialized:
              self.initialize()
  
          if isinstance(filename_or_checkpoint, str):
              checkpoint = torch.load(filename_or_checkpoint, map_location=self.device)
          # if state dict comes from nn.DataParallel but we use non-parallel model here then the state dict keys do not
          # match. Use heuristic to make it match
          new_state_dict = {}
          for k, value in checkpoint["network_weights"].items():
              key = k
              if key not in self.network.state_dict().keys() and key.startswith("module."):
                  key = key[7:]
              new_state_dict[key] = value
  
          self.my_init_kwargs = checkpoint["init_args"]
          self.current_epoch = checkpoint["current_epoch"]
          self.logger.load_checkpoint(checkpoint["logging"])
          self._best_ema = checkpoint["_best_ema"]
          self.inference_allowed_mirroring_axes = (
              checkpoint["inference_allowed_mirroring_axes"]
              if "inference_allowed_mirroring_axes" in checkpoint.keys()
              else self.inference_allowed_mirroring_axes
          )
  
          # messing with state dict naming schemes. Facepalm.
          if self.is_ddp:
              if isinstance(self.network.module, OptimizedModule):
                  self.network.module._orig_mod.load_state_dict(new_state_dict)
              else:
                  self.network.module.load_state_dict(new_state_dict)
          else:
              if isinstance(self.network, OptimizedModule):
                  self.network._orig_mod.load_state_dict(new_state_dict)
              else:
                  self.network.load_state_dict(new_state_dict)
  
          # it's fine to do this every time we load because configure_optimizers will be a no-op if the correct optimizer
          # and lr scheduler are already set up
          if self.current_epoch < self.warmup_duration_whole_net:
              self.optimizer, self.lr_scheduler = self.configure_optimizers("warmup_all")
          else:
              self.optimizer, self.lr_scheduler = self.configure_optimizers("train")
  
          self.optimizer.load_state_dict(checkpoint["optimizer_state"])
          if self.grad_scaler is not None:
              if checkpoint["grad_scaler_state"] is not None:
                  self.grad_scaler.load_state_dict(checkpoint["grad_scaler_state"])