Method train_and_evaluate

orbit/controller.py:353–398

Runs interleaved training and evaluation. This method interleaves calls to `self.train()` and `self.evaluate()`, training the model until the global step count equals `train_steps`, and running an evaluation for `eval_steps` every `eval_interval` training steps. In addition, this me

(
      self,
      train_steps: int,
      eval_steps: int = -1,
      eval_interval: Optional[int] = None,
  )

    return eval_output

  def train_and_evaluate(
      self,
      train_steps: int,
      eval_steps: int = -1,
      eval_interval: Optional[int] = None,
  ) -> Optional[runner.Output]:
    """Runs interleaved training and evaluation.

    This method interleaves calls to `self.train()` and `self.evaluate()`,
    training the model until the global step count equals `train_steps`, and
    running an evaluation for `eval_steps` every `eval_interval` training steps.
    In addition, this method will run a final evaluation at the end of the
    training sequence.

    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:
      train_steps: The global step count to train up to.
      eval_steps: The number of steps to run during an evaluation. If -1, this
        method will evaluate over the entire evaluation dataset.
      eval_interval: The number of training steps to run between evaluations. If
        set, training will always stop every `eval_interval` steps, even if this
        results in a shorter inner loop than specified by `steps_per_loop`
        setting. If None, evaluation will only be performed after training is
        complete.

    Returns:
      The evaluation results as a dictionary mapping names to NumPy values.
    """
    self._require("trainer", for_method="train_and_evaluate")
    self._require("evaluator", for_method="train_and_evaluate")

    output = None
    current_step = self.global_step.numpy()  # Cache, since this is expensive.
    eval_interval = eval_interval or (train_steps - current_step)
    while current_step < train_steps:
      interval = min(train_steps - current_step, eval_interval)
      num_steps = current_step + interval
      self.train(steps=num_steps, checkpoint_at_completion=False)
      output = self.evaluate(steps=eval_steps)
      current_step = self.global_step.numpy()
    self._maybe_save_checkpoint(check_interval=False)
    self._sync_on_async_checkpointing()
    return output

  def evaluate_continuously(
      self,