Method train

orbit/controller.py:257–289

Runs training until the specified global step count has been reached. This method makes calls to `self.trainer.train()` until the global step count is equal to `steps`. It will additionally save checkpoints (if a `CheckpointManager` was passed to `Controller.__init__`) and summarize

(self, steps: int, checkpoint_at_completion: bool = True)

    _orbit_api_gauge.get_cell().set(True)

  def train(self, steps: int, checkpoint_at_completion: bool = True):
    """Runs training until the specified global step count has been reached.

    This method makes calls to `self.trainer.train()` until the global step
    count is equal to `steps`. It will additionally save checkpoints (if a
    `CheckpointManager` was passed to `Controller.__init__`) and summarize
    training output (if `summary_dir` is set).

    When async checkpointing is enabled, a sync is triggered at the end of this
    method to make sure any ongoing async checkpoint saving is finished before
    returning.

    Args:
      steps: The global step count to train up to.
      checkpoint_at_completion: Whether to save a checkpoint when this method
        returns (regardless of the checkpointing interval). Defaults to `True`.
    """
    self._require("trainer", for_method="train")

    # TODO(momernick): Support steps=None or -1 (training to exhaustion).
    current_step = self.global_step.numpy()  # Cache, since this is expensive.
    _log(f"train | step: {current_step: 6d} | training until step {steps}...")
    while current_step < steps:
      # Calculates steps to run for the next train loop.
      num_steps = min(steps - current_step, self.steps_per_loop)
      self._train_n_steps(num_steps)
      self._maybe_save_checkpoint()
      current_step = self.global_step.numpy()

    if checkpoint_at_completion:
      self._maybe_save_checkpoint(check_interval=False)

    self._sync_on_async_checkpointing()

  def evaluate(self, steps: int = -1) -> Optional[runner.Output]:
    """Runs evaluation for the given number of steps.