(self, task, total_loss, n_steps, elapsed_time,
                             optimizer_metrics, summary_writer)
Logs training related metrics.

    Logs:
     * current learning rate,
     * steps per second,
     * average training loss,
     * average metrics returned from the optimizer
    to the provided summary writer. Training loss is also logged to stdout.

    Args:
      task: Current training task.
      total_loss: Total training loss accumulated over n_steps training steps.
      n_steps: Number of steps over which the metrics were accumulated.
      elapsed_time: Time of execution of n_steps training steps.
      optimizer_metrics: Dict from optimizer metric name to metric values.
      summary_writer: Jaxboard summary writer for saving provided metrics.
Logs training related metrics.
Logs training related metrics .
def _log_training_progress(self, task, total_loss, n_steps, elapsed_time,
                             optimizer_metrics, summary_writer):
    """Logs training related metrics.

    Logs:
     * current learning rate,
     * steps per second,
     * average training loss,
     * average metrics returned from the optimizer
    to the provided summary writer. Training loss is also logged to stdout.

    Args:
      task: Current training task.
      total_loss: Total training loss accumulated over n_steps training steps.
      n_steps: Number of steps over which the metrics were accumulated.
      elapsed_time: Time of execution of n_steps training steps.
      optimizer_metrics: Dict from optimizer metric name to metric values.
      summary_writer: Jaxboard summary writer for saving provided metrics.
    """
    # only here do avoid potential divide-by-0
    n_steps = max(1, n_steps)
    _log('')  # Separator for visibility on terminals.
    if self.step == 1:
      self._log_n_weights()
    self._log_step('Ran %d train steps in %0.2f secs' % (n_steps, elapsed_time))
    self.log_summary(
        {task.loss_name: total_loss / float(n_steps)},
        summary_writer, 'metrics/', 'train')
    if self.step == 1:
      self._save_gin(summary_writer)
    train_parameters = {
        'learning_rate': task.learning_rate(self.step),
        'steps per second': n_steps / elapsed_time,
    }
    # Average optimizer_metrics over n_steps.
    optimizer_metrics = {k: v / n_steps for k, v in optimizer_metrics.items()}
    train_parameters.update(optimizer_metrics)
    self.log_summary(
        train_parameters, summary_writer, 'training/', 'train', stdout=False)
def _log_training_progress ( self , task , total_loss , n_steps , elapsed_time , optimizer_metrics , summary_writer ) : # only here do avoid potential divide-by-0 n_steps = max ( 1 , n_steps ) _log ( '' ) # Separator for visibility on terminals. if self . step == 1 : self . _log_n_weights ( ) self . _log_step ( 'Ran %d train steps in %0.2f secs' % ( n_steps , elapsed_time ) ) self . log_summary ( { task . loss_name : total_loss / float ( n_steps ) } , summary_writer , 'metrics/' , 'train' ) if self . step == 1 : self . _save_gin ( summary_writer ) train_parameters = { 'learning_rate' : task . learning_rate ( self . step ) , 'steps per second' : n_steps / elapsed_time , } # Average optimizer_metrics over n_steps. optimizer_metrics = { k : v / n_steps for k , v in optimizer_metrics . items ( ) } train_parameters . update ( optimizer_metrics ) self . log_summary ( train_parameters , summary_writer , 'training/' , 'train' , stdout = False )