content: # Copyright (c) Open-MMLab. All rights reserved.
  import os.path as osp
  import platform
  import shutil
  import time
  import warnings
  
  import torch
  
  import mmcv
  from mmcv.runner import EpochBasedRunner
  
  class AnnealingRunner(EpochBasedRunner):
  
      def run_iter(self, data_batch, train_mode, **kwargs):
          
          if 'annealing' in kwargs:
              kwargs.update(epoch=self.epoch)
              kwargs.update(total_epoch=self.max_epochs)
              kwargs.update(iter=self._iter)
          
          if self.batch_processor is not None:
              outputs = self.batch_processor(
                  self.model, data_batch, train_mode=train_mode, **kwargs)
          elif train_mode:
              outputs = self.model.train_step(data_batch, self.optimizer,
                                              **kwargs)
          else:
              outputs = self.model.val_step(data_batch, self.optimizer, **kwargs)
          if not isinstance(outputs, dict):
              raise TypeError('"batch_processor()" or "model.train_step()"'
                              'and "model.val_step()" must return a dict')
          if 'log_vars' in outputs:
              self.log_buffer.update(outputs['log_vars'], outputs['num_samples'])
          self.outputs = outputs
  
      def train(self, data_loader, **kwargs):
          self.model.train()
          self.mode = 'train'
          self.data_loader = data_loader
          self._max_iters = self._max_epochs * len(self.data_loader)
          self.call_hook('before_train_epoch')
          time.sleep(2)  # Prevent possible deadlock during epoch transition
          for i, data_batch in enumerate(self.data_loader):
              self._inner_iter = i
              self.call_hook('before_train_iter')
              self.run_iter(data_batch, train_mode=True, **kwargs)
              self.call_hook('after_train_iter')
              self._iter += 1
  
          self.call_hook('after_train_epoch')
          self._epoch += 1
  
      def val(self, data_loader, **kwargs):
          self.model.eval()
          self.mode = 'val'
          self.data_loader = data_loader
          self.call_hook('before_val_epoch')
          time.sleep(2)  # Prevent possible deadlock during epoch transition
          for i, data_batch in enumerate(self.data_loader):
              self._inner_iter = i
              self.call_hook('before_val_iter')
              with torch.no_grad():
                  self.run_iter(data_batch, train_mode=False, **kwargs)
              self.call_hook('after_val_iter')
  
          self.call_hook('after_val_epoch')