text: # Copyright 2021 RangiLyu.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  import copy
  import os
  import warnings
  import json
  import torch
  import logging
  
  from pytorch_lightning import LightningModule
  from typing import Any, List
  from nanodet.util import mkdir, gather_results
  
  from ..model.arch import build_model
  
  
  class TrainingTask(LightningModule):
      """
      Pytorch Lightning module of a general training task.
      Including training, evaluating and testing.
      Args:
          cfg: Training configurations
          evaluator: Evaluator for evaluating the model performance.
      """
  
      def __init__(self, cfg, evaluator=None):
          super(TrainingTask, self).__init__()
          self.cfg = cfg
          self.model = build_model(cfg.model)
          self.evaluator = evaluator
          self.save_flag = -10
          self.log_style = 'NanoDet'  # Log style. Choose between 'NanoDet' or 'Lightning'
          # TODO: use callback to log
  
      def forward(self, x):
          x = self.model(x)
          return x
  
      @torch.no_grad()
      def predict(self, batch, batch_idx=None, dataloader_idx=None):
          preds = self.forward(batch['img'])
          results = self.model.head.post_process(preds, batch)
          return results
  
      def on_train_start(self) -> None:
          self.lr_scheduler.last_epoch = self.current_epoch-1
  
      def training_step(self, batch, batch_idx):
          preds, loss, loss_states = self.model.forward_train(batch)
  
          # log train losses
          if self.log_style == 'Lightning':
              self.log('lr', self.optimizers().param_groups[0]['lr'], on_step=True, on_epoch=False, prog_bar=True)
              for k, v in loss_states.items():
                  self.log('Train/'+k, v, on_step=True, on_epoch=True, prog_bar=True, sync_dist=True)
          elif self.log_style == 'NanoDet' and self.global_step % self.cfg.log.interval == 0:
              lr = self.optimizers().param_groups[0]['lr']
              log_msg = 'Train|Epoch{}/{}|Iter{}({})| lr:{:.2e}| '.format(self.current_epoch+1,
                  self.cfg.schedule.total_epochs, self.global_step, batch_idx, lr)
              self.scalar_summary('Train_loss/lr', 'Train', lr, self.global_step)
              for l in loss_states:
                  log_msg += '{}:{:.4f}| '.format(l, loss_states[l].mean().item())
                  self.scalar_summary('Train_loss/' + l, 'Train', loss_states[l].mean().item(), self.global_step)
              self.info(log_msg)
  
          return loss
  
      def training_epoch_end(self, outputs: List[Any]) -> None:
          self.trainer.save_checkpoint(os.path.join(self.cfg.save_dir, 'model_last.ckpt'))
          self.lr_scheduler.step()
  
      def validation_step(self, batch, batch_idx):
          preds, loss, loss_states = self.model.forward_train(batch)
  
          if self.log_style == 'Lightning':
              self.log('Val/loss', loss, on_step=True, on_epoch=False, prog_bar=True, logger=False)
              for k, v in loss_states.items():
                  self.log('Val/' + k, v, on_step=False, on_epoch=True, prog_bar=False, sync_dist=True)
          elif self.log_style == 'NanoDet' and batch_idx % self.cfg.log.interval == 0:
              lr = self.optimizers().param_groups[0]['lr']
              log_msg = 'Val|Epoch{}/{}|Iter{}({})| lr:{:.2e}| '.format(self.current_epoch+1,
                  self.cfg.schedule.total_epochs, self.global_step, batch_idx, lr)
              for l in loss_states:
                  log_msg += '{}:{:.4f}| '.format(l, loss_states[l].mean().item())
              self.info(log_msg)
  
          dets = self.model.head.post_process(preds, batch)
          return dets
  
      def validation_epoch_end(self, validation_step_outputs):
          """
          Called at the end of the validation epoch with the outputs of all validation steps.
          Evaluating results and save best model.
          Args:
              validation_step_outputs: A list of val outputs
  
          """
          results = {}
          for res in validation_step_outputs:
              results.update(res)
          all_results = gather_results(results)
          if all_results:
              eval_results = self.evaluator.evaluate(all_results, self.cfg.save_dir, rank=self.local_rank)
              metric = eval_results[self.cfg.evaluator.save_key]
              # save best model
              if metric > self.save_flag:
                  self.save_flag = metric
                  best_save_path = os.path.join(self.cfg.save_dir, 'model_best')
                  mkdir(self.local_rank, best_save_path)
                  self.trainer.save_checkpoint(os.path.join(best_save_path, "model_best.ckpt"))
                  txt_path = os.path.join(best_save_path, "eval_results.txt")
                  if self.local_rank < 1:
                      with open(txt_path, "a") as f:
                          f.write("Epoch:{}\n".format(self.current_epoch+1))
                          for k, v in eval_results.items():
                              f.write("{}: {}\n".format(k, v))
              else:
                  warnings.warn('Warning! Save_key is not in eval results! Only save model last!')
              if self.log_style == 'Lightning':
                  for k, v in eval_results.items():
                      self.log('Val_metrics/' + k, v, on_step=False, on_epoch=True, prog_bar=False, sync_dist=True)
              elif self.log_style == 'NanoDet':
                  for k, v in eval_results.items():
                      self.scalar_summary('Val_metrics/' + k, 'Val', v, self.current_epoch+1)
          else:
              self.info('Skip val on rank {}'.format(self.local_rank))
  
      def test_step(self, batch, batch_idx):
          dets = self.predict(batch, batch_idx)
          return dets
  
      def test_epoch_end(self, test_step_outputs):
          results = {}
          for res in test_step_outputs:
              results.update(res)
          all_results = gather_results(results)
          if all_results:
              res_json = self.evaluator.results2json(all_results)
              json_path = os.path.join(self.cfg.save_dir, 'results.json')
              json.dump(res_json, open(json_path, 'w'))
  
              if self.cfg.test_mode == 'val':
                  eval_results = self.evaluator.evaluate(all_results, self.cfg.save_dir, rank=self.local_rank)
                  txt_path = os.path.join(self.cfg.save_dir, "eval_results.txt")
                  with open(txt_path, "a") as f:
                      for k, v in eval_results.items():
                          f.write("{}: {}\n".format(k, v))
          else:
              self.info('Skip test on rank {}'.format(self.local_rank))
  
      def configure_optimizers(self):
          """
          Prepare optimizer and learning-rate scheduler
          to use in optimization.
  
          Returns:
              optimizer
          """
          optimizer_cfg = copy.deepcopy(self.cfg.schedule.optimizer)
          name = optimizer_cfg.pop('name')
          build_optimizer = getattr(torch.optim, name)
          optimizer = build_optimizer(params=self.parameters(), **optimizer_cfg)
  
          schedule_cfg = copy.deepcopy(self.cfg.schedule.lr_schedule)
          name = schedule_cfg.pop('name')
          build_scheduler = getattr(torch.optim.lr_scheduler, name)
          self.lr_scheduler = build_scheduler(optimizer=optimizer, **schedule_cfg)
          # lr_scheduler = {'scheduler': self.lr_scheduler,
          #                 'interval': 'epoch',
          #                 'frequency': 1}
          # return [optimizer], [lr_scheduler]
  
          return optimizer
  
      def optimizer_step(self,
                         epoch=None,
                         batch_idx=None,
                         optimizer=None,
                         optimizer_idx=None,
                         optimizer_closure=None,
                         on_tpu=None,
                         using_native_amp=None,
                         using_lbfgs=None):
          """
          Performs a single optimization step (parameter update).
          Args:
              epoch: Current epoch
              batch_idx: Index of current batch
              optimizer: A PyTorch optimizer
              optimizer_idx: If you used multiple optimizers this indexes into that list.
              optimizer_closure: closure for all optimizers
              on_tpu: true if TPU backward is required
              using_native_amp: True if using native amp
              using_lbfgs: True if the matching optimizer is lbfgs
          """
          # warm up lr
          if self.trainer.global_step <= self.cfg.schedule.warmup.steps:
              if self.cfg.schedule.warmup.name == 'constant':
                  warmup_lr = self.cfg.schedule.optimizer.lr * self.cfg.schedule.warmup.ratio
              elif self.cfg.schedule.warmup.name == 'linear':
                  k = (1 - self.trainer.global_step / self.cfg.schedule.warmup.steps) * (1 - self.cfg.schedule.warmup.ratio)
                  warmup_lr = self.cfg.schedule.optimizer.lr * (1 - k)
              elif self.cfg.schedule.warmup.name == 'exp':
                  k = self.cfg.schedule.warmup.ratio ** (1 - self.trainer.global_step / self.cfg.schedule.warmup.steps)
                  warmup_lr = self.cfg.schedule.optimizer.lr * k
              else:
                  raise Exception('Unsupported warm up type!')
              for pg in optimizer.param_groups:
                  pg['lr'] = warmup_lr
  
          # update params
          optimizer.step(closure=optimizer_closure)
          optimizer.zero_grad()
  
      def get_progress_bar_dict(self):
          # don't show the version number
          items = super().get_progress_bar_dict()
          items.pop("v_num", None)
          items.pop("loss", None)
          return items
  
      def scalar_summary(self, tag, phase, value, step):
          """
          Write Tensorboard scalar summary log.
          Args:
              tag: Name for the tag
              phase: 'Train' or 'Val'
              value: Value to record
              step: Step value to record
  
          """
          if self.local_rank < 1:
              self.logger.experiment.add_scalars(tag, {phase: value}, step)
  
      def info(self, string):
          if self.local_rank < 1:
              logging.info(string)
  
  
  
  
  
  
  
  
  