official/core/base_trainer.py

# Copyright 2023 The TensorFlow 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.

"""Standard Trainer implementation.

The base trainer implements the Orbit `StandardTrainable` and
`StandardEvaluable` interfaces. Trainers inside this project should be
interchangable and independent on model architectures and tasks.
"""
import functools
from typing import Union, Optional
from absl import logging
import gin
import orbit
import tensorflow as tf, tf_keras

from official.core import base_task
from official.core import config_definitions
from official.modeling import optimization

ExperimentConfig = config_definitions.ExperimentConfig
TrainerConfig = config_definitions.TrainerConfig


class _AsyncTrainer(orbit.StandardTrainer, orbit.StandardEvaluator):
  """Trainer class for both sync and async Strategy."""

  def init_async(self):
    """Initializes the Async Trainer base class."""
    assert isinstance(self._strategy, tf.distribute.Strategy)
    self._is_async = isinstance(
        self._strategy, tf.distribute.experimental.ParameterServerStrategy)
    self._coordinator = None
    if self._is_async:
      self._coordinator = (
          tf.distribute.experimental.coordinator.ClusterCoordinator(
              self._strategy))

  def coordinator_for_async(
      self,
  ) -> tf.distribute.experimental.coordinator.ClusterCoordinator:
    if not self._coordinator:
      raise ValueError(
          "Coordinator uninitialized for async run. Call init_async() first."
      )
    return self._coordinator

  def join(self):
    """Join all async steps. Only useful in aysnc training."""
    if getattr(self, "_is_async", False):
      self.coordinator_for_async().join()

  def create_train_loop_fn(self):
    """Creates a eval loop from the given step function and options."""
    train_loop_fn = super().create_train_loop_fn()
    if getattr(self, "_is_async", False):

      def _async_loop_fn(iterator, num_steps):
        self.coordinator_for_async().schedule(
            train_loop_fn, args=(iterator, num_steps)
        )

      return _async_loop_fn
    else:
      return train_loop_fn

  def create_eval_loop_fn(self, has_state: bool):
    """Creates a training loop from the given step function and options."""
    eval_loop_fn = super().create_eval_loop_fn(has_state)

    if getattr(self, "_is_async", False):
      if has_state:
        raise ValueError(
            "Stateful eval loop is not supported in async training.")

      def _async_loop_fn(iterator, num_steps, state=None, reduce_fn=None):
        assert state is None
        assert reduce_fn is None
        self.coordinator_for_async().schedule(
            eval_loop_fn, args=(iterator, num_steps)
        )

      return _async_loop_fn
    else:
      return eval_loop_fn

  def distribute_dataset(self, dataset_or_fn, *args, **kwargs):
    """A utility function to help create a `tf.distribute.DistributedDataset`.

    Args:
      dataset_or_fn: A instance of `tf.data.Dataset`, or a "dataset function"
        returning a `tf.data.Dataset`. If it is a function, it may optionally
        have an argument named `input_context` which will be passed a
        `tf.distribute.InputContext` instance.
      *args: Any positional arguments to pass through to `dataset_or_fn`.
      **kwargs: Any keyword arguments to pass through to `dataset_or_fn`.

    Returns:
      A distributed Dataset.
    """
    if getattr(self, "_is_async", False):
      per_worker_dataset_fn = functools.partial(
          orbit.utils.make_distributed_dataset, self._strategy, dataset_or_fn,
          *args, **kwargs)
      per_worker_dataset_fn = tf.function(per_worker_dataset_fn)

      return self.coordinator_for_async().create_per_worker_dataset(
          per_worker_dataset_fn
      )
    else:
      return orbit.utils.make_distributed_dataset(self._strategy, dataset_or_fn,
                                                  *args, **kwargs)


def get_runtime_options(config: ExperimentConfig):
  """Get tf.distribute.RunOptions from config."""
  xla_options = {}
  if config.runtime.tpu_enable_xla_dynamic_padder is not None:
    xla_options["enable_xla_dynamic_padder"] = (
        config.runtime.tpu_enable_xla_dynamic_padder)
  return tf.distribute.RunOptions(
      experimental_xla_options=tf.tpu.XLAOptions(**xla_options))


@gin.configurable
class Trainer(_AsyncTrainer):
  """Implements the common trainer shared for TensorFlow models."""

  # pylint: disable=super-init-not-called
  def __init__(
      self,
      config: ExperimentConfig,
      task: base_task.Task,
      model: tf_keras.Model,
      optimizer: tf.optimizers.Optimizer,
      train: bool = True,
      evaluate: bool = True,
      train_dataset: Optional[Union[tf.data.Dataset,
                                    tf.distribute.DistributedDataset]] = None,
      validation_dataset: Optional[Union[
          tf.data.Dataset, tf.distribute.DistributedDataset]] = None,
      checkpoint_exporter=None):
    """Initialize common trainer for TensorFlow models.

    Args:
      config: An `ExperimentConfig` instance specifying experiment config.
      task: A base_task.Task instance.
      model: The model instance, e.g. a tf_keras.Model instance.
      optimizer: tf.optimizers.Optimizer instance.
      train: bool, whether or not this trainer will be used for training.
        default to True.
      evaluate: bool, whether or not this trainer will be used for evaluation.
        default to True.
      train_dataset: a dataset object created for training. With tf.distribute,
        it needs to be a `DistributedDataset`.
      validation_dataset: a dataset object created for evaluation. With
        tf.distribute, it needs to be a `DistributedDataset`. The evaluator will
        create a dataset iterator for each eval round, so the dataset does not
        need to repeat.
      checkpoint_exporter: an object that has the `maybe_export_checkpoint`
        interface.
    """
    # Gets the current distribution strategy. If not inside any strategy scope,
    # it gets a single-replica no-op strategy.
    self._strategy = tf.distribute.get_strategy()
    self._validate_params(
        config,
        check_train_data=train_dataset is None,
        check_validation_data=validation_dataset is None)
    self._config = config
    self._task = task
    self._model = model
    self._optimizer = optimizer
    self._checkpoint_exporter = checkpoint_exporter
    self._recovery = None
    # Runtime options are only applied to train_step.
    # We use default for eval_step.
    self._runtime_options = get_runtime_options(config)

    # Creates a shadow copy of the weights to store weights moving average.
    if isinstance(self._optimizer, optimization.ExponentialMovingAverage
                 ) and not self._optimizer.has_shadow_copy:
      self._optimizer.shadow_copy(self._model)

    # global_step increases by 1 after each training iteration.
    # We should have global_step.numpy() == self.optimizer.iterations.numpy()
    # when there is only 1 optimizer.
    self._global_step = orbit.utils.create_global_step()
    if hasattr(self.model, "checkpoint_items"):
      checkpoint_items = self.model.checkpoint_items
    else:
      checkpoint_items = {}
    self._checkpoint = tf.train.Checkpoint(
        global_step=self.global_step,
        model=self.model,
        optimizer=self.optimizer,
        **checkpoint_items)

    self._train_loss = tf_keras.metrics.Mean("training_loss", dtype=tf.float32)
    self._validation_loss = tf_keras.metrics.Mean(
        "validation_loss", dtype=tf.float32)
    model_metrics = model.metrics if hasattr(model, "metrics") else []

    self.init_async()

    if train:
      self._train_metrics = self.task.build_metrics(
          training=True) + model_metrics
      train_dataset = train_dataset or self.distribute_dataset(
          self.task.build_inputs, self.config.task.train_data)
      orbit.StandardTrainer.__init__(
          self,
          train_dataset,
          options=orbit.StandardTrainerOptions(
              use_tf_while_loop=config.trainer.train_tf_while_loop,
              use_tf_function=config.trainer.train_tf_function,
              use_tpu_summary_optimization=config.trainer.allow_tpu_summary))

    if evaluate:
      self._validation_metrics = self.task.build_metrics(
          training=False) + model_metrics
      validation_dataset = validation_dataset or self.distribute_dataset(
          self.task.build_inputs, self.config.task.validation_data)
      orbit.StandardEvaluator.__init__(
          self,
          validation_dataset,
          options=orbit.StandardEvaluatorOptions(
              use_tf_function=config.trainer.eval_tf_function,
              use_tf_while_loop=config.trainer.eval_tf_while_loop))

  def _validate_params(self,
                       config,
                       check_train_data=True,
                       check_validation_data=True):
    r"""Validates if the configuration object passed to the Trainer.

    The experiment configuration should be structured as:
    \trainer
    \task
      \train_data
      \validation_data

    Args:
      config: a namedtuple, dataclass, ConfigDict, etc.
      check_train_data: whether to check task.train_data field.
      check_validation_data: whether to check task.validation_data field.
    """
    if not hasattr(config, "trainer"):
      raise AttributeError("The trainer requires the configuration contains an"
                           " attribute `trainer`.")

    if not hasattr(config, "task"):
      raise AttributeError("The trainer requires the configuration contains an"
                           " attribute `task`.")

    if check_train_data and not hasattr(config.task, "train_data"):
      raise AttributeError("The trainer requires the configuration contains an"
                           " attribute `task.train_data`.")

    if check_validation_data and not hasattr(config.task, "validation_data"):
      raise AttributeError("The trainer requires the configuration contains an"
                           " attribute `task.validation_data`.")

  @property
  def strategy(self):
    return self._strategy

  @property
  def config(self):
    return self._config

  @property
  def task(self):
    return self._task

  @property
  def model(self):
    return self._model

  @property
  def optimizer(self):
    if hasattr(self, "_optimizer"):
      return self._optimizer
    else:
      return None

  @property
  def global_step(self):
    return self._global_step

  @property
  def train_loss(self):
    """Accesses the training loss metric object."""
    return self._train_loss

  @property
  def validation_loss(self):
    """Accesses the validation loss metric object."""
    return self._validation_loss

  @property
  def train_metrics(self):
    """Accesses all training metric objects."""
    return self._train_metrics

  @property
  def validation_metrics(self):
    """Accesses all validation metric metric objects."""
    return self._validation_metrics

  def initialize(self):
    """A callback function.

    This function will be called when no checkpoint found for the model.
    If there is a checkpoint, the checkpoint will be loaded and this function
    will not be called. Tasks may use this callback function to load a
    pretrained checkpoint, saved under a directory other than the model_dir.
    """
    self.task.initialize(self.model)

  @property
  def checkpoint(self):
    """Accesses the training checkpoint."""
    return self._checkpoint

  @property
  def checkpoint_exporter(self):
    """Accesses the checkpoint exporter."""
    return self._checkpoint_exporter

  def train_loop_end(self):
    """See base class."""
    self.join()
    logs = {}
    for metric in self.train_metrics + [self.train_loss]:
      logs[metric.name] = metric.result()
      metric.reset_states()
    if callable(self.optimizer.learning_rate):
      # Maybe a self-implemented optimizer does not have `optimizer.iterations`.
      # So just to be safe here.
      if hasattr(self.optimizer, "iterations"):
        logs["learning_rate"] = self.optimizer.learning_rate(
            self.optimizer.iterations)
      else:
        logs["learning_rate"] = self.optimizer.learning_rate(self.global_step)
    else:
      logs["learning_rate"] = self.optimizer.learning_rate
    return logs

  def next_train_inputs(self, iterator):
    """Fetches the next inputs for the model during train.

    This method consumes the input iterator and returns the next inputs for the
    model.

    This method provides a way to control how to fetch the next model input, and
    what data to send to the model.

    Note: This function runs on the host side when accelerators are used.

    Note: Depending on the training setup this may or may not run in eager mode.
    In most cases it will be run in graph mode.

    Args:
      iterator: Dataset iterator to generate the next inputs from.

    Returns:
      The inputs to the model.
    """
    return next(iterator)

  def train_step(self, iterator):
    """See base class."""

    def step_fn(inputs):
      if self.config.runtime.enable_xla and (self.config.runtime.num_gpus > 0):
        task_train_step = tf.function(self.task.train_step, jit_compile=True)
      else:
        task_train_step = self.task.train_step
      logs = task_train_step(
          inputs,
          model=self.model,
          optimizer=self.optimizer,
          metrics=self.train_metrics)
      self._train_loss.update_state(logs[self.task.loss])
      self.global_step.assign_add(1)

    inputs = self.next_train_inputs(iterator)
    self.strategy.run(step_fn, args=(inputs,), options=self._runtime_options)

  def eval_begin(self):
    """Sets up metrics."""
    for metric in self.validation_metrics + [self.validation_loss]:
      metric.reset_states()
    # Swaps weights to test on weights moving average.
    if self.optimizer and isinstance(self.optimizer,
                                     optimization.ExponentialMovingAverage):
      self.optimizer.swap_weights()

  def next_eval_inputs(self, iterator):
    """Fetches the next inputs for the model during eval.

    This method consumes the input iterator and returns the next inputs for the
    model and an additional logs dict. The output dict remains in the host (not
    sent to GPUs/TPUs) and is merged with the model outputs which will be
    processed later in `aggregate_logs`. This is useful for sending extra logs
    downstream that are not compatible with the accelerators.

    Note: This function runs on the host side when accelerators are used.

    Note: Depending on the training setup this may or may not run in eager mode.
    In most cases it will be run in graph mode.

    Args:
      iterator: Dataset iterator to generate the next inputs from.

    Returns:
      The inputs to the model, and an additional logs dictionnary. The logs
      are not passed to the model, instead they are merged with model output
      logs.
    """
    passthrough_logs = dict()
    return next(iterator), passthrough_logs

  def eval_step(self, iterator):
    """See base class."""

    def step_fn(inputs):
      logs = self.task.validation_step(
          inputs, model=self.model, metrics=self.validation_metrics)
      if self.task.loss in logs:
        self._validation_loss.update_state(logs[self.task.loss])
      return logs

    inputs, passthrough_logs = self.next_eval_inputs(iterator)
    distributed_outputs = self.strategy.run(step_fn, args=(inputs,))
    logs = tf.nest.map_structure(
        self.strategy.experimental_local_results, distributed_outputs
    )

    if set(logs.keys()) & set(passthrough_logs.keys()):
      logging.warning(
          (
              "Conflict between the pasthrough log keys and the returned model"
              " log keys. Found %r keys in the passthrough logs and %r keys in"
              " the model logs. Model log keys takes precedence."
          ),
          logs.keys(),
          passthrough_logs.keys(),
      )

    return passthrough_logs | logs

  def eval_end(self, aggregated_logs=None):
    """Processes evaluation results."""
    self.join()
    logs = {}
    for metric in self.validation_metrics:
      logs[metric.name] = metric.result()
    if self.validation_loss.count.numpy() != 0:
      logs[self.validation_loss.name] = self.validation_loss.result()
    else:
      # `self.validation_loss` metric was not updated, because the validation
      # loss was not returned from the task's `validation_step` method.
      logging.info("The task did not report validation loss.")
    if aggregated_logs:
      metrics = self.task.reduce_aggregated_logs(
          aggregated_logs, global_step=self.global_step)
      logs.update(metrics)

    if self._checkpoint_exporter:
      self._checkpoint_exporter.maybe_export_checkpoint(
          self.checkpoint, logs, self.global_step.numpy())
      metric_name = self.config.trainer.best_checkpoint_eval_metric
      logs["best_" +
           metric_name] = self._checkpoint_exporter.best_ckpt_logs[metric_name]

    # Swaps back weights after testing when EMA is used.
    # This happens after best checkpoint export so that average weights used for
    # eval are exported instead of regular weights.
    if self.optimizer and isinstance(self.optimizer,
                                     optimization.ExponentialMovingAverage):
      self.optimizer.swap_weights()
    return logs

  def eval_reduce(self, state=None, step_outputs=None):
    return self.task.aggregate_logs(state, step_outputs)