controller.py

# Lint as: python3
# Copyright 2020 The Orbit Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""A light weight utilities to train TF2 models."""

import time
from typing import Callable, Dict, Optional, Text, Union
from absl import logging
import numpy as np
from orbit import runner
from orbit import utils

import tensorflow as tf


def _log_info(message: Text):
  """Logs `message` to the `info` log, and also prints to stdout."""
  logging.info(message)
  print(message)


class Controller:
  """Class that facilitates training and evaluation of models."""

  def __init__(
      self,
      strategy: Optional[tf.distribute.Strategy] = None,
      trainer: Optional[runner.AbstractTrainer] = None,
      evaluator: Optional[runner.AbstractEvaluator] = None,
      global_step: Optional[tf.Variable] = None,
      # Train related
      steps_per_loop: Optional[int] = None,
      checkpoint_manager: Optional[tf.train.CheckpointManager] = None,
      # Summary related
      summary_interval: Optional[int] = None,
      summary_dir: Optional[Text] = None,
      # Evaluation related
      eval_summary_dir: Optional[Text] = None):
    """Constructs a `Controller` instance.

    Args:
      strategy: An instance of `tf.distribute.Strategy`.
      trainer: An instance of `orbit.AbstractTrainer`, which represents model
        training details.
      evaluator: An instance of `orbit.AbstractEvaluator`, which represents
        model evaluation details.
      global_step: An integer `tf.Variable` indicating the global training step
        number. Usually this can be obtained from `iterations` property of the
        model's optimizer (e.g. `self.optimizer.iterations`), or users can
        create their own global step variable as well. If the users create their
        own global step variable, it is recommended to create the `tf.Variable`
        inside strategy scope, and with
        `aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA`.
      steps_per_loop: The number of steps to run in each "inner loop" of
        training (passed to the `num_steps` parameter of `trainer.train`).
      checkpoint_manager: An instance of `tf.train.CheckpointManager`.
      summary_interval: Step interval for training summaries. Note that this
        argument only applies to the summaries inside `trainer.train` function.
        Summaries outside like "steps_per_second" and outputs from
        `trainer.train` function will always be enabled. If set, the value
        should be divisible by steps_per_loop.
      summary_dir: The directory to restore and write checkpoints and summaries.
        If None, it will be set to `checkpoint_manager.directory`.
      eval_summary_dir: The directory to write eval summaries. If None, it will
        be set to `summary_dir`.

    Raises:
      ValueError: If both `trainer` and `evaluator` are None.
      ValueError: If `steps_per_loop` is not a positive integer.
      ValueError: If `summary_interval` is not a positive integer or it cannot
        be divisible by `steps_per_loop`.
    """
    if trainer is None and evaluator is None:
      raise ValueError("`trainer` and `evaluator` should not both be None")

    if trainer is not None:
      if steps_per_loop is None:
        raise ValueError("`steps_per_loop` is required when `trainer` is "
                         "provided.")

      if not isinstance(steps_per_loop, int) or steps_per_loop < 1:
        raise ValueError("`steps_per_loop` should be a positive integer")

      if summary_interval is not None:
        if summary_interval <= 0:
          raise ValueError("`summary_interval` should be larger than 0")
        if summary_interval % steps_per_loop != 0:
          raise ValueError("The summary interval ({}) must be a multiple "
                           "of the steps_per_loop ({})".format(
                               summary_interval, steps_per_loop))

    self.trainer = trainer
    self.evaluator = evaluator

    self.strategy = strategy or tf.distribute.get_strategy()

    self.global_step = global_step
    self.checkpoint_manager = checkpoint_manager

    if summary_dir is None and checkpoint_manager:
      summary_dir = checkpoint_manager.directory

    if self.trainer is not None:
      self.step_timer = None
      self.steps_per_loop = steps_per_loop
      self.summary_interval = summary_interval
      self.summary_manager = utils.SummaryManager(
          summary_dir, tf.summary.scalar, global_step=self.global_step)

    eval_summary_writer = None
    if self.evaluator is not None:
      eval_summary_dir = eval_summary_dir or summary_dir
      if eval_summary_dir == summary_dir and self.trainer is not None:
        # Reuse the summary writer if train and evaluation summary directory
        # are the same.
        self.eval_summary_manager = self.summary_manager
      else:
        self.eval_summary_manager = utils.SummaryManager(
            eval_summary_dir, tf.summary.scalar, global_step=self.global_step)

    if self.global_step is not None:
      tf.summary.experimental.set_step(self.global_step)

    # Restores the model if needed.
    # TODO(momernick): We probably only want to do this on certain occasions?
    if self.checkpoint_manager is not None:
      checkpoint_interval = self.checkpoint_manager.checkpoint_interval
      model_restored = self.restore_checkpoint()
      if not model_restored and (checkpoint_interval and
                                 self.trainer is not None):
        # If the model is not restored from a checkpoint, and
        # `checkpoint_interval` is enabled for training, save an initial
        # checkpoint.
        self.save_checkpoint()

  def train(self, steps: int, checkpoint_at_completion: bool = True):
    """Runs training.

    This method calls the `train` method on the Trainable object until the
    global step count is equal to `steps`. It will optionally save checkpoints,
    if a CheckpointManager was passed to the Controller instance's `__init__`.

    Args:
      steps: The global step count to train up to.
      checkpoint_at_completion: Whether to save a checkpoint when this method
        returns. Defaults to True (write the checkpoint). This is always
        triggered, regardless of the checkpointing interval.
    """
    if self.trainer is None:
      raise ValueError("`self.trainer` is required when calling `train` "
                       "method.")
    if self.global_step is None:
      raise ValueError("`self.global_step` is required when calling `train` "
                       "method.")

    # TODO(momernick): Support steps=None or -1 (training to exhaustion).
    current_step = self.global_step.numpy()  # This is an expensive access.
    while current_step < steps:
      logging.info("Train at step %s of %s", 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()  # This is an expensive access.

    if checkpoint_at_completion:
      self.save_checkpoint()

  def evaluate(self, steps: int = None) -> Optional[Dict[Text, np.number]]:
    """Runs evaluation.

    This method calls the `evaluate` method on the Evaluator object for `steps`
    steps, then writes the returned summaries (if any).

    Args:
      steps: The number of steps to evaluate for.

    Returns:
      The evaluation results as a dictionary of numpy values.

    Raises:
      ValueError: If no checkpoint found in `self.checkpoint_manager.directory`.
      ValueError: If `evaluator` is not provided.
    """
    if self.evaluator is None:
      raise ValueError("`evaluator` must be provided to call `evaluate()` "
                       "method.")

    steps = steps or -1
    current_step = self.global_step.numpy()
    if steps > 0:
      logging.info("Running %s steps of evaluation at train step: %s", steps,
                   current_step)
      steps = tf.convert_to_tensor(steps, dtype=tf.int32)
    else:
      logging.info("Evaluating at train step: %s", current_step)

    with self.eval_summary_manager.summary_writer.as_default():
      eval_outputs = self.evaluator.evaluate(steps)

    if eval_outputs:
      eval_outputs = tf.nest.map_structure(utils.get_value, eval_outputs)

    info = "step: {}        evaluation metric: {}".format(
        current_step, eval_outputs)
    _log_info(info)

    self.eval_summary_manager.write_summaries(eval_outputs)
    self.eval_summary_manager.flush()

    return eval_outputs

  def restore_checkpoint(self, checkpoint_path: Text = None):
    """Restore or initialize the model.

    Args:
      checkpoint_path: An optional string indicates the checkpoint path to
        restore. If None, will restore from `self.checkpoint_manager`.

    Returns:
      The path to the restored checkpoint if a restore happened, or None
        if no restore occurred.
    """
    with self.strategy.scope():
      # Checkpoint restoring should be inside scope. b/139450638
      if checkpoint_path is not None:
        self.checkpoint_manager.checkpoint.restore(checkpoint_path)
        return checkpoint_path
      return self.checkpoint_manager.restore_or_initialize()

  def save_checkpoint(self):
    """Checkpoint the model.

    This method will write a checkpoint containing the current state of the
    model.

    Raises:
      ValueError: if no CheckpointManager was provided to this Controller's
        init args.
    """
    self._maybe_save_checkpoint(force_trigger=True)

  def train_and_evaluate(self,
                         train_steps: int = None,
                         eval_steps: int = None,
                         eval_interval: int = None):
    """Train and evaluate in an interleaved manner.

    This method will train the model until the global step count equals
    `train_steps`, 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.

    Args:
      train_steps: The global step count to train up to.
      eval_steps: The number of steps to run during an evaluation. If None,
        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.

    Raises:
      ValueError: If eval_interval is not a multiple of self.steps_per_loop.
    """
    current_step = self.global_step.numpy()  # This is an expensive access.
    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)
      self.evaluate(steps=eval_steps)
      current_step = self.global_step.numpy()  # This is an expensive access.
    self.save_checkpoint()

  def evaluate_continuously(self,
                            steps: int = None,
                            timeout: Optional[Union[int, float]] = None,
                            timeout_fn: Optional[Callable[[], bool]] = None):
    """Monitor a directory and evaluate on checkpoints in it.

    This method continuously monitors a directory as specified by this
    Controller's CheckpointManager init arg and runs evaluation on the
    checkpoints found there.

    Args:
      steps: The number of steps to run when evaluating.
      timeout: The maximum number of seconds to wait between checkpoints. See
        tf.train.checkpoints_iterator documentation.
      timeout_fn: Optional callable to call after a timeout. If the function
        returns True, then it means that no new checkpoints will be generated
        and the iterator will exit.

    Raises:
      ValueError: If no checkpoint found in `self.checkpoint_manager.directory`.
      ValueError: If `evaluator` was not provided as a controller init arg.

    """
    for checkpoint_path in tf.train.checkpoints_iterator(
        self.checkpoint_manager.directory,
        timeout=timeout,
        timeout_fn=timeout_fn):
      self.restore_checkpoint(checkpoint_path)
      self.evaluate(steps)

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

    It will also write training outputs to summaries if there is any.

    Args:
      num_steps: An integer indicates how many steps to run for this training
        loop.

    Raises:
      RuntimeError: If `global_step` is not updated correctly in
        `trainer.train`.
    """
    if not self.step_timer:
      self.step_timer = StepTimer(self.global_step)

    # Calculates steps to run for the next train loop.
    current_step = self.global_step.numpy()
    logging.info("Entering training loop at step %s to run %s steps",
                 current_step, num_steps)
    current_step += num_steps
    num_steps = tf.convert_to_tensor(num_steps, dtype=tf.int32)

    with self.summary_manager.summary_writer.as_default():
      # Create a lambda that returns true when summaries should be written.
      should_record = False  # Allows static optimization in no-summary cases.
      if self.summary_interval:
        should_record = lambda: (self.global_step % self.summary_interval == 0)
      with tf.summary.record_if(should_record):
        train_outputs = self.trainer.train(num_steps)

    # Updates and verifies the current step after a training loop finishes.
    if current_step != self.global_step.numpy():
      raise RuntimeError("`trainer.train` function is not updating "
                         "`global_step` correctly, expected: %s, actual: %s" %
                         (current_step, self.global_step.numpy()))

    # Print information like metrics and steps_per_second after a training
    # loop.
    if train_outputs:
      train_outputs = tf.nest.map_structure(utils.get_value, train_outputs)

    train_outputs = train_outputs or {}
    steps_per_second = self.step_timer.steps_per_second()
    info = "step: {}        steps_per_second: {:.2f}        {}".format(
        current_step, steps_per_second, train_outputs)
    _log_info(info)

    train_outputs["steps_per_second"] = steps_per_second
    self.summary_manager.write_summaries(train_outputs)

  def _maybe_save_checkpoint(self, force_trigger: bool = False):
    """Save checkpoints if necessary.

    Args:
      force_trigger: A boolean indicates whether to force saving checkpoints
        regardless of the checkpoint interval.

    Returns:
      A boolean indicating whether a checkpoint was saved.
    """
    if self.checkpoint_manager and self.checkpoint_manager.checkpoint_interval:
      ckpt_path = self.checkpoint_manager.save(
          checkpoint_number=self.global_step.numpy(),
          check_interval=not force_trigger)
      if ckpt_path is not None:
        logging.info("Saved checkpoints in %s", ckpt_path)
        return True
    return False


class StepTimer:
  """Utility class for measuring steps/second."""

  def __init__(self, step):
    self.step = step
    self.start()

  def start(self):
    self.last_iteration = self.step.numpy()
    self.last_time = time.time()

  def steps_per_second(self, restart=True):
    value = ((self.step.numpy() - self.last_iteration) /
             (time.time() - self.last_time))
    if restart:
      self.start()
    return value