class_docstr: Train/eval directpose
class_docstr_tok: ["Train\/eval","directpose"]
code: def train_step(self):
      cfg = self.cfg
      writer = self.writer
      self.model.train()
      end = time.perf_counter()
      data = next(self.dataloader)
      self.base_iter += 1
      if self.base_iter >= self.max_iter:
          return
  
      data_time = time.perf_counter() - end
  
      loss_dict = self.model(data)
      losses = sum(loss_dict.values())
      _detect_anomaly(losses, loss_dict, self.base_iter)
  
      metrics_dict = loss_dict
      metrics_dict["data_time"] = data_time
  
      self.optimizer.zero_grad()
      losses.backward()
      self.optimizer.step()
  
      batch_time = time.perf_counter() - end
      end = time.perf_counter()
      metrics_dict["batch_time"] = batch_time
  
      # gather all metrics
      metrics_dict = {
          k: v.detach().cpu().item()
          if isinstance(v, torch.Tensor)
          else float(v)
          for k, v in metrics_dict.items()
      }
      all_metrics_dict = comm.gather(metrics_dict)
      if (
          self.base_iter % cfg.CONFIG.LOG.DISPLAY_FREQ == 0
          and cfg.DDP_CONFIG.GPU_WORLD_RANK == 0
      ):
          eta_str = None
          if "data_time" in all_metrics_dict[0]:
              # data_time among workers can have high variance. The actual latency
              # caused by data_time is the maximum among workers.
              data_time = np.max(
                  [x.pop("data_time") for x in all_metrics_dict]
              )
          if "batch_time" in all_metrics_dict[0]:
              # batch_time among workers can have high variance. The actual latency
              # caused by batch_time is the maximum among workers.
              batch_time = np.max(
                  [x.pop("batch_time") for x in all_metrics_dict]
              )
              self.iter_timer.update(batch_time)
              eta = (
                  self.max_iter - self.base_iter
              ) * self.iter_timer.avg
              eta_str = str(datetime.timedelta(seconds=int(eta)))
          # average the rest metrics
          metrics_dict = {
              k: np.mean([x[k] for x in all_metrics_dict])
              for k in all_metrics_dict[0].keys()
          }
          total_losses_reduced = sum(
              loss for loss in metrics_dict.values()
          )
          for param in self.optimizer.param_groups:
              lr = param["lr"]
          print_string = "Iter: [{0}/{1}]".format(
              self.base_iter, self.max_iter
          )
          if eta_str is not None:
              print_string += f" ETA: {eta_str} "
          print_string += " data_time: {data_time:.3f}, batch time: {batch_time:.3f}".format(
              data_time=data_time, batch_time=batch_time
          )
          print_string += " loss: {loss:.5f}".format(
              loss=total_losses_reduced
          )
          print_string += " lr: {lr:.6f}".format(lr=lr)
          iteration = self.base_iter
          writer.add_scalar(
              "total_loss", total_losses_reduced, iteration
          )
          writer.add_scalar("learning_rate", lr, iteration)
          if len(metrics_dict) > 1:
              for km, kv in metrics_dict.items():
                  writer.add_scalar(km, kv, iteration)
                  print_string += f" {km}: {kv:.2f}"
          logger.info(print_string)
      self.scheduler.step()
  
file_path: gluon-cv/gluoncv/torch/utils/task_utils/pose.py
imports: import time
  import logging
  import datetime
  import numpy
  import torch
  from ..None import comm
  from ..utils import AverageMeter
  from ..model_utils import save_model
  from ..optimizer import maybe_add_gradient_clipping
  from ..eval_utils.coco_eval import COCOEvaluator
label_desc: Use image node_id 2 for calling the DirectposePipeline obj's underlying member method code with example usage: obj.train_step() without return types