Method _train_n_steps

orbit/controller.py:491–541

Runs training for `num_steps` steps. Also prints/logs updates about training progress, and summarizes training output (if output is returned from `self.trainer.train()`, and if `self.summary_dir` is set). Args: num_steps: An integer specifying how many steps of training to ru

(self, num_steps: int)

    return self._steps_per_loop

  def _train_n_steps(self, num_steps: int):
    """Runs training for `num_steps` steps.

    Also prints/logs updates about training progress, and summarizes training
    output (if output is returned from `self.trainer.train()`, and if
    `self.summary_dir` is set).

    Args:
      num_steps: An integer specifying how many steps of training to run.

    Raises:
      RuntimeError: If `global_step` is not properly incremented by `num_steps`
        after calling `self.trainer.train(num_steps)`.
    """
    if not self.step_timer:
      self.step_timer = StepTimer(self.global_step)
    current_step = self.global_step.numpy()

    with self.summary_manager.summary_writer().as_default():
      should_record = False  # Allows static optimization in no-summary cases.
      if self.summary_interval:
        # Create a predicate to determine when summaries should be written.
        should_record = lambda: (self.global_step % self.summary_interval == 0)
      assert isinstance(self.trainer, runner.AbstractTrainer)
      with tf.summary.record_if(should_record):
        num_steps_tensor = tf.convert_to_tensor(num_steps, dtype=tf.int32)
        train_output = self.trainer.train(num_steps_tensor)

    # Verify that global_step was updated properly, then update current_step.
    expected_step = current_step + num_steps
    if self.global_step.numpy() != expected_step:
      message = (
          f"`trainer.train({num_steps})` did not update `global_step` by "
          f"{num_steps}. Old value was {current_step}, expected updated value "
          f"to be {expected_step}, but it was {self.global_step.numpy()}.")
      logging.warning(message)

    train_output = train_output or {}
    for action in self.train_actions:
      action(train_output)
    train_output = tf.nest.map_structure(utils.get_value, train_output)

    current_step = self.global_step.numpy()
    steps_per_second = self.step_timer.steps_per_second()
    _log(f"train | step: {current_step: 6d} | "
         f"steps/sec: {steps_per_second: 6.1f} | "
         f"output: {_format_output(train_output)}")

    train_output["steps_per_second"] = steps_per_second
    self.summary_manager.write_summaries(train_output)
    self.summary_manager.flush()

  def _maybe_save_checkpoint(self, check_interval: bool = True):
    """Conditionally saves a checkpoint.