class_docstr: Hook for writing the losses, learning rates and meters to `tensorboard <https
  ://www.tensorflow.org/tensorboard`>_.
  
  Global steps are counted in terms of the number of samples processed.
class_docstr_tok: ["Hook","for","writing","the","losses",",","learning","rates","and","meters","to","`","tensorboard","<","https",":","\/\/www.tensorflow.org\/tensorboard","`",">","_",".","Global","steps","are","counted","in","terms","of","the","number","of","samples","processed","."]
docstr: Add the losses and learning rates to tensorboard.
docstr_tok: ["Add","the","losses","and","learning","rates","to","tensorboard","."]
code: def on_phase_end(self, task) -> None:
      """Add the losses and learning rates to tensorboard."""
      if self.learning_rates is None:
          logging.warning("learning_rates is not initialized")
          return
  
      phase_type = task.phase_type
      cum_sample_fetch_times = self._get_cum_sample_fetch_times(
          phase_type
      )
  
      batches = len(task.losses)
      if batches == 0 or not is_primary():
          return
  
      phase_type_idx = (
          task.train_phase_idx if task.train else task.eval_phase_idx
      )
  
      logging.info(
          f"Plotting to Tensorboard for {phase_type} phase {phase_type_idx}"
      )
  
      for i in range(0, len(self.wall_times), self.log_period):
          global_step = (
              i
              + self.state.step_count[phase_type]
              - len(self.wall_times)
              + 1
          )
          if cum_sample_fetch_times:
              self.tb_writer.add_scalar(
                  f"Speed/{phase_type}/cumulative_sample_fetch_time",
                  cum_sample_fetch_times[i],
                  global_step=global_step,
                  walltime=self.wall_times[i],
              )
          if task.train:
              self.tb_writer.add_scalar(
                  "Learning Rate/train",
                  self.learning_rates[i],
                  global_step=global_step,
                  walltime=self.wall_times[i],
              )
  
      if task.train:
          for name, parameter in task.base_model.named_parameters():
              self.tb_writer.add_histogram(
                  f"Parameters/{name}",
                  parameter,
                  global_step=phase_type_idx,
              )
  
      if torch.cuda.is_available():
          self.tb_writer.add_scalar(
              f"Memory/{phase_type}/peak_allocated",
              torch.cuda.max_memory_allocated(),
              global_step=phase_type_idx,
          )
  
      loss_avg = sum(task.losses) / batches
  
      loss_key = f"Losses/{phase_type}"
      self.tb_writer.add_scalar(
          loss_key, loss_avg, global_step=phase_type_idx
      )
  
      # plot meters which return a dict
      for meter in task.meters:
          if not isinstance(meter.value, dict):
              log.warn(
                  f"Skipping meter {meter.name} with value: {meter.value}"
              )
              continue
          for name, value in meter.value.items():
              if isinstance(value, float):
                  meter_key = f"Meters/{phase_type}/{meter.name}/{name}"
                  self.tb_writer.add_scalar(
                      meter_key, value, global_step=phase_type_idx
                  )
              else:
                  log.warn(
                      f"Skipping meter name {meter.name}/{name} with value: {value}"
                  )
                  continue
  
      if hasattr(task, "perf_log"):
          for perf in task.perf_log:
              phase_idx = perf["phase_idx"]
              tag = perf["tag"]
              for metric_name, metric_value in perf.items():
                  if metric_name in ["phase_idx", "tag"]:
                      continue
  
                  self.tb_writer.add_scalar(
                      f"Speed/{tag}/{metric_name}",
                      metric_value,
                      global_step=phase_idx,
                  )
  
      # flush so that the plots aren't lost if training crashes soon after
      self.tb_writer.flush()
      logging.info("Done plotting to Tensorboard")
  
file_path: ClassyVision/classy_vision/hooks/tensorboard_plot_hook.py
imports: import logging
  import time
  from itertools import accumulate
  from typing import Any, Dict, List, Optional, Tuple
  import torch
  from classy_vision.generic.distributed_util import all_reduce_max, is_primary
  from classy_vision.hooks import register_hook
  from classy_vision.hooks.classy_hook import ClassyHook
label_desc: Use image node_id 7 for calling the TensorboardPlotHook obj's underlying member method code with example usage: obj.on_phase_end(task) without return types