Method evaluate

orbit/controller.py:291–351

Runs evaluation for the given number of steps. This method calls `self.evaluator.evaluate(steps)`, then writes the returned summaries (if any). Args: steps: The number of evaluation steps to run. The value `-1` is reserved as a special sentinel to indicate a "complete" ev

(self, steps: int = -1)

    self._sync_on_async_checkpointing()

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

    This method calls `self.evaluator.evaluate(steps)`, then writes the returned
    summaries (if any).

    Args:
      steps: The number of evaluation steps to run. The value `-1` is reserved
        as a special sentinel to indicate a "complete" evaluation that runs
        until the underlying dataset is exhausted. Support for this is dependent
        on the specific `evaluator` being used.

    Returns:
      The evaluation results as a dictionary mapping names to NumPy values.

    Raises:
      ValueError: If `evaluator` was not provided to `Controller.__init__`.
      ValueError: If no checkpoint is present in `checkpoint_manager.directory`.
      ValueError: If `steps` is not a positive value or -1.
    """
    self._require("evaluator", for_method="evaluate")

    if steps > 0:
      steps_msg = f"running {steps} steps of evaluation..."
    elif steps == -1:
      steps_msg = "running complete evaluation..."
    else:
      raise ValueError(f"`steps` ({steps}) should be > 0, or == -1.")

    current_step = self.global_step.numpy()
    _log(f" eval | step: {current_step: 6d} | {steps_msg}")

    start = time.time()
    assert isinstance(self.evaluator, runner.AbstractEvaluator)
    with self.eval_summary_manager.summary_writer().as_default():
      steps_tensor = tf.convert_to_tensor(steps, dtype=tf.int32)
      eval_output = self.evaluator.evaluate(steps_tensor)
    elapsed = time.time() - start

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

    if steps > 0:
      # Only log if steps has been specified.
      steps_per_second = steps / elapsed
      eval_output["steps_per_second"] = steps_per_second
      steps_per_second_log = f"steps/sec: {steps_per_second: 6.1f} | "
    else:
      steps_per_second_log = ""

    _log(f" eval | step: {current_step: 6d} | "
         f"{steps_per_second_log}"
         f"eval time: {elapsed: 6.1f} sec | "
         f"output: {_format_output(eval_output)}")

    self.eval_summary_manager.write_summaries(eval_output)