text: """
  Dyhead Training Script.
  
  This script is a simplified version of the training script in detectron2/tools.
  """
  import os
  import sys
  import itertools
  import logging
  import time
  # fmt: off
  sys.path.insert(1, os.path.join(sys.path[0], '..'))
  # fmt: on
  
  from typing import Any, Dict, List, Set
  
  import torch
  from torch.nn.parallel import DistributedDataParallel
  
  import detectron2.utils.comm as comm
  
  from detectron2.checkpoint import DetectionCheckpointer
  from detectron2.config import get_cfg
  from detectron2.data import build_detection_train_loader, build_detection_test_loader
  from detectron2.data.samplers import TrainingSampler
  from detectron2.data.build import get_detection_dataset_dicts
  from detectron2.engine import SimpleTrainer, DefaultTrainer, default_argument_parser, default_setup, launch
  from detectron2.evaluation import COCOEvaluator, verify_results
  from detectron2.solver.build import maybe_add_gradient_clipping
  from detectron2.utils.logger import setup_logger
  
  from dyhead import add_dyhead_config
  from extra import add_extra_config
  
  
  class Trainer(DefaultTrainer):
      """
      Extension of the Trainer class adapted to DETR.
      """
  
      def __init__(self, cfg):
          """
          Args:
              cfg (CfgNode):
          """
          super(DefaultTrainer, self).__init__()
          logger = logging.getLogger("detectron2")
          if not logger.isEnabledFor(logging.INFO):  # setup_logger is not called for d2
              setup_logger()
          cfg = DefaultTrainer.auto_scale_workers(cfg, comm.get_world_size())
  
          model = self.build_model(cfg)
          optimizer = self.build_optimizer(cfg, model)
          data_loader = self.build_train_loader(cfg)
          # Build DDP Model with find_unused_parameters to add flexibility.
          if comm.get_world_size() > 1:
              model = DistributedDataParallel(
                  model, device_ids=[comm.get_local_rank()], broadcast_buffers=False,
                  find_unused_parameters=True
              )
          self._trainer = SimpleTrainer(model, data_loader, optimizer)
          self.scheduler = self.build_lr_scheduler(cfg, optimizer)
          self.checkpointer = DetectionCheckpointer(
              # Assume you want to save checkpoints together with logs/statistics
              model,
              cfg.OUTPUT_DIR,
              optimizer=optimizer,
              scheduler=self.scheduler,
          )
          self.start_iter = 0
          self.max_iter = cfg.SOLVER.MAX_ITER
          self.cfg = cfg
  
          self.register_hooks(self.build_hooks())
  
      @classmethod
      def build_evaluator(cls, cfg, dataset_name, output_folder=None):
          """
          Create evaluator(s) for a given dataset.
          This uses the special metadata "evaluator_type" associated with each builtin dataset.
          For your own dataset, you can simply create an evaluator manually in your
          script and do not have to worry about the hacky if-else logic here.
          """
          if output_folder is None:
              output_folder = os.path.join(cfg.OUTPUT_DIR, "inference")
          return COCOEvaluator(dataset_name, cfg, True, output_folder)
  
      @classmethod
      def build_train_loader(cls, cfg):
          dataset = get_detection_dataset_dicts(
              cfg.DATASETS.TRAIN,
              filter_empty=cfg.DATALOADER.FILTER_EMPTY_ANNOTATIONS,
              min_keypoints=cfg.MODEL.ROI_KEYPOINT_HEAD.MIN_KEYPOINTS_PER_IMAGE
              if cfg.MODEL.KEYPOINT_ON
              else 0,
              proposal_files=cfg.DATASETS.PROPOSAL_FILES_TRAIN if cfg.MODEL.LOAD_PROPOSALS else None,
          )
  
          mapper = None
          if cfg.SEED!=-1:
              sampler = TrainingSampler(len(dataset), seed=cfg.SEED)
          else:
              sampler = None
          return build_detection_train_loader(cfg, dataset=dataset, mapper=mapper, sampler=sampler)
  
      @classmethod
      def build_test_loader(cls, cfg, dataset_name):
          mapper = None
          return build_detection_test_loader(cfg, mapper=mapper, dataset_name=dataset_name)
  
      @classmethod
      def build_optimizer(cls, cfg, model):
          params: List[Dict[str, Any]] = []
          memo: Set[torch.nn.parameter.Parameter] = set()
          for key, value in model.named_parameters(recurse=True):
              if not value.requires_grad:
                  continue
              # Avoid duplicating parameters
              if value in memo:
                  continue
              memo.add(value)
              lr = cfg.SOLVER.BASE_LR
              weight_decay = cfg.SOLVER.WEIGHT_DECAY
              if "bias" in key:
                  lr = cfg.SOLVER.BASE_LR * cfg.SOLVER.BIAS_LR_FACTOR
                  weight_decay = cfg.SOLVER.WEIGHT_DECAY_BIAS
              params += [{"params": [value], "lr": lr, "weight_decay": weight_decay}]
  
          def maybe_add_full_model_gradient_clipping(optim):  # optim: the optimizer class
              # detectron2 doesn't have full model gradient clipping now
              clip_norm_val = cfg.SOLVER.CLIP_GRADIENTS.CLIP_VALUE
              enable = (
                  cfg.SOLVER.CLIP_GRADIENTS.ENABLED
                  and cfg.SOLVER.CLIP_GRADIENTS.CLIP_TYPE == "full_model"
                  and clip_norm_val > 0.0
              )
  
              class FullModelGradientClippingOptimizer(optim):
                  def step(self, closure=None):
                      all_params = itertools.chain(*[x["params"] for x in self.param_groups])
                      torch.nn.utils.clip_grad_norm_(all_params, clip_norm_val)
                      super().step(closure=closure)
  
              return FullModelGradientClippingOptimizer if enable else optim
  
          optimizer_type = cfg.SOLVER.OPTIMIZER
          if optimizer_type == "SGD":
              optimizer = maybe_add_full_model_gradient_clipping(torch.optim.SGD)(
                  params, cfg.SOLVER.BASE_LR, momentum=cfg.SOLVER.MOMENTUM
              )
          elif optimizer_type == "ADAMW":
              optimizer = maybe_add_full_model_gradient_clipping(torch.optim.AdamW)(
                  params, cfg.SOLVER.BASE_LR
              )
          else:
              raise NotImplementedError(f"no optimizer type {optimizer_type}")
          if not cfg.SOLVER.CLIP_GRADIENTS.CLIP_TYPE == "full_model":
              optimizer = maybe_add_gradient_clipping(cfg, optimizer)
          return optimizer
  
      def resume_or_load(self, resume=True):
          checkpoint = self.checkpointer.resume_or_load(self.cfg.MODEL.WEIGHTS, resume=resume)
          if resume and self.checkpointer.has_checkpoint():
              self.start_iter = checkpoint.get("iteration", -1) + 1
  
  
  def setup(args):
      """
      Create configs and perform basic setups.
      """
      cfg = get_cfg()
      add_dyhead_config(cfg)
      add_extra_config(cfg)
      cfg.merge_from_file(args.config_file)
      cfg.merge_from_list(args.opts)
      cfg.freeze()
      default_setup(cfg, args)
      return cfg
  
  
  def main(args):
      cfg = setup(args)
  
      if args.eval_only:
          model = Trainer.build_model(cfg)
          DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load(cfg.MODEL.WEIGHTS, resume=args.resume)
          res = Trainer.test(cfg, model)
          if comm.is_main_process():
              verify_results(cfg, res)
          return res
  
      trainer = Trainer(cfg)
      trainer.resume_or_load(resume=True)
      return trainer.train()
  
  
  if __name__ == "__main__":
      args = default_argument_parser().parse_args()
      print("Command Line Args:", args)
      launch(
          main,
          args.num_gpus,
          num_machines=args.num_machines,
          machine_rank=int(os.environ['AZ_BATCHAI_TASK_INDEX']) \
                          if 'AZ_BATCHAI_TASK_INDEX' in os.environ else args.machine_rank,
          dist_url="tcp://"+os.environ['AZ_BATCH_MASTER_NODE'] \
                          if 'AZ_BATCH_MASTER_NODE' in os.environ else args.dist_url,
          args=(args,),
      )
  