Method __init__

orbit/controller.py:94–255

Initializes a `Controller` instance. Note that if `checkpoint_manager` is provided and there are checkpoints in the associated model directory, the model will be restored from the most recent checkpoint during this `__init__` method. Args: global_step: An integer `tf.Variable

(
      self,
      *,  # Makes all args keyword only.
      global_step: tf.Variable,
      trainer: Optional[runner.AbstractTrainer] = None,
      evaluator: Optional[runner.AbstractEvaluator] = None,
      strategy: Optional[tf.distribute.Strategy] = None,
      # Actions
      train_actions: Optional[Iterable[Action]] = None,
      eval_actions: Optional[Iterable[Action]] = None,
      # Train related
      steps_per_loop: Optional[Union[int, Callable[[int], int]]] = None,
      checkpoint_manager: Optional[tf.train.CheckpointManager] = None,
      enable_async_checkpointing: bool = False,
      # Summary related
      summary_interval: Optional[int] = None,
      summary_dir: Optional[str] = None,
      # Evaluation related
      eval_summary_dir: Optional[str] = None,
      summary_manager: Optional[utils.SummaryManagerInterface] = None,
      eval_summary_manager: Optional[utils.SummaryManagerInterface] = None)

  """

  def __init__(
      self,
      *,  # Makes all args keyword only.
      global_step: tf.Variable,
      trainer: Optional[runner.AbstractTrainer] = None,
      evaluator: Optional[runner.AbstractEvaluator] = None,
      strategy: Optional[tf.distribute.Strategy] = None,
      # Actions
      train_actions: Optional[Iterable[Action]] = None,
      eval_actions: Optional[Iterable[Action]] = None,
      # Train related
      steps_per_loop: Optional[Union[int, Callable[[int], int]]] = None,
      checkpoint_manager: Optional[tf.train.CheckpointManager] = None,
      enable_async_checkpointing: bool = False,
      # Summary related
      summary_interval: Optional[int] = None,
      summary_dir: Optional[str] = None,
      # Evaluation related
      eval_summary_dir: Optional[str] = None,
      summary_manager: Optional[utils.SummaryManagerInterface] = None,
      eval_summary_manager: Optional[utils.SummaryManagerInterface] = None):
    """Initializes a `Controller` instance.

    Note that if `checkpoint_manager` is provided and there are checkpoints in
    the associated model directory, the model will be restored from the most
    recent checkpoint during this `__init__` method.

    Args:
      global_step: An integer `tf.Variable` storing the global training step
        number. Usually this can be obtained from the `iterations` property of
        the model's optimizer (e.g. `trainer.optimizer.iterations`). In cases
        where multiple optimizers are used, or if one model "step" corresponds
        to more than one update to model parameters, users can create and
        increment their own global step variable as well. In this case it is
        recommended to create the `tf.Variable` inside the distribution strategy
        scope, with `aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA` (see
        also `orbit.utils.create_global_step()`).
      trainer: An instance of `orbit.AbstractTrainer`, which implements the
        inner training loop.
      evaluator: An instance of `orbit.AbstractEvaluator`, which implements
        evaluation.
      strategy: An instance of `tf.distribute.Strategy`. If not provided, the
        strategy will be initialized from the current in-scope strategy using
        `tf.distribute.get_strategy()`.
      train_actions: Optional `orbit.Action`s to call after each block of
        `steps_per_loop` training steps are run. These will be called with the
        output of `trainer.train`.
      eval_actions: Optional `orbit.Action`s to call after each evaluation.
        These will be called with the output of `evaluator.evaluate`.
      steps_per_loop: Optional integer to indicate the number of steps to run in
        each inner loop of training (passed as the `num_steps` parameter of
        `trainer.train`). It can be also a callable which takes the current
        global step value as input and returns the number of steps to run as
        output.
      checkpoint_manager: An instance of `tf.train.CheckpointManager`. If
        provided and there are checkpoints in the associated model directory,
        the model will be restored from the most recent checkpoint inside this
        `__init__` method. If not provided, the `Controller` will not