(
      self,
      model,
      tasks,
      eval_model=None,
      eval_tasks=None,
      output_dir=None,
      checkpoint_at=None,
      checkpoint_low_metric=None,
      checkpoint_high_metric=None,
      permanent_checkpoint_at=None,
      eval_at=None,
      which_task=None,
      n_devices=None,
      random_seed=None,
      loss_chunk_size=0,
      use_memory_efficient_trainer=False,
      adasum=False,
      callbacks=None,
  )
Configures a training ``Loop``, including a random initialization.

    Args:
      model: Trax layer, representing the core model to be trained. Loss
          functions and eval functions (a.k.a. metrics) are considered to be
          outside the core model, taking core model output and data labels as
          their two inputs.
      tasks: List of :py:class:`TrainTask` instances, which define the training
          data, loss function, and optimizer to be used in respective tasks in
          this training loop. It can also be a single :py:class:`TrainTask`
          instance which is treated in the same way as a singleton list.
      eval_model: Optional Trax layer, representing model used for evaluation,
          e.g., with dropout turned off. If ``None``, the training model (model)
          will be used.
      eval_tasks: List of :py:class:`EvalTask` instances which define how to
          evaluate the model: which validation data to use and which metrics to
          report. Evaluation on each of the tasks and will run and be reported
          separately which allows to score a model on different subtasks. This
          argument can also be ``None``, in which case no evals will be run, or
          a single :py:class:`EvalTask`, which wil be treated in the same way
          as a singleton list.
      output_dir: Path telling where to save outputs (evals and checkpoints).
          Can be ``None`` if both ``eval_task`` and ``checkpoint_at`` are
          ``None``.
      checkpoint_at: Function (integer --> boolean) telling, for step n, whether
          that step should have its checkpoint saved. If ``None``, the default
          is periodic checkpointing at ``task.n_steps_per_checkpoint``.
      checkpoint_low_metric: Name of metric, or None. The metric name must
          be one of the metric names from the evals in ``eval_tasks``. At
          checkpoint times determined by ``checkpoint_at``, a separate
          specially named checkpoint will be saved (overwriting any previous
          version) if the designated metric reaches a value less than or equal
          to any previous recorded low value. No such checkpoint is saved if
          arg value is `None`.
      checkpoint_high_metric: Name of metric, or None. The metric name must
          be one of the metric names from the evals in ``eval_tasks``. At
          checkpoint times determined by ``checkpoint_at``, a separate
          specially named checkpoint will be saved (overwriting any previous
          version) if the designated metric reaches a value greater than or
          equal to any previous recorded high value. No such checkpoint is
          saved if arg value is `None`.
      permanent_checkpoint_at: Function (integer --> boolean) telling,
          for step n, whether that step should have its checkpoint saved
          permanently. If ``None``, the default is periodic checkpointing at
          ``task.n_steps_per_permanent_checkpoint``.
      eval_at: Function (integer --> boolean) that says, for training step n,
          whether that step should run evals. If ``None``, run evals on the
          first step and on every N'th step, as determined by the first
          training task.
      which_task: Function (integer --> integer) indicating which task should be
          used at which training step. Can be set to ``None`` in single-task
          training.
      n_devices: integer or ``None``, the number of devices for this
          computation.
      random_seed: the random seed to use; time/os dependent if ``None``
          (default).
      loss_chunk_size: int, if > 0 use chunks of this size to make loss
        computation more more memory-efficient.
      use_memory_efficient_trainer: whether to use a special memory-efficient
        trainer; if set to 2, the memory efficiency if very aggressive
      adasum: if True, use adaptive summation for multi-device gradients
      callbacks: List of subclasses of StepCallback to call on training
        steps.
Configures a training ``Loop``, including a random initialization.
Configures a training Loop including a random initialization .
def __init__(
      self,
      model,
      tasks,
      eval_model=None,
      eval_tasks=None,
      output_dir=None,
      checkpoint_at=None,
      checkpoint_low_metric=None,
      checkpoint_high_metric=None,
      permanent_checkpoint_at=None,
      eval_at=None,
      which_task=None,
      n_devices=None,
      random_seed=None,
      loss_chunk_size=0,
      use_memory_efficient_trainer=False,
      adasum=False,
      callbacks=None,
  ):
    """Configures a training ``Loop``, including a random initialization.

    Args:
      model: Trax layer, representing the core model to be trained. Loss
          functions and eval functions (a.k.a. metrics) are considered to be
          outside the core model, taking core model output and data labels as
          their two inputs.
      tasks: List of :py:class:`TrainTask` instances, which define the training
          data, loss function, and optimizer to be used in respective tasks in
          this training loop. It can also be a single :py:class:`TrainTask`
          instance which is treated in the same way as a singleton list.
      eval_model: Optional Trax layer, representing model used for evaluation,
          e.g., with dropout turned off. If ``None``, the training model (model)
          will be used.
      eval_tasks: List of :py:class:`EvalTask` instances which define how to
          evaluate the model: which validation data to use and which metrics to
          report. Evaluation on each of the tasks and will run and be reported
          separately which allows to score a model on different subtasks. This
          argument can also be ``None``, in which case no evals will be run, or
          a single :py:class:`EvalTask`, which wil be treated in the same way
          as a singleton list.
      output_dir: Path telling where to save outputs (evals and checkpoints).
          Can be ``None`` if both ``eval_task`` and ``checkpoint_at`` are
          ``None``.
      checkpoint_at: Function (integer --> boolean) telling, for step n, whether
          that step should have its checkpoint saved. If ``None``, the default
          is periodic checkpointing at ``task.n_steps_per_checkpoint``.
      checkpoint_low_metric: Name of metric, or None. The metric name must
          be one of the metric names from the evals in ``eval_tasks``. At
          checkpoint times determined by ``checkpoint_at``, a separate
          specially named checkpoint will be saved (overwriting any previous
          version) if the designated metric reaches a value less than or equal
          to any previous recorded low value. No such checkpoint is saved if
          arg value is `None`.
      checkpoint_high_metric: Name of metric, or None. The metric name must
          be one of the metric names from the evals in ``eval_tasks``. At
          checkpoint times determined by ``checkpoint_at``, a separate
          specially named checkpoint will be saved (overwriting any previous
          version) if the designated metric reaches a value greater than or
          equal to any previous recorded high value. No such checkpoint is
          saved if arg value is `None`.
      permanent_checkpoint_at: Function (integer --> boolean) telling,
          for step n, whether that step should have its checkpoint saved
          permanently. If ``None``, the default is periodic checkpointing at
          ``task.n_steps_per_permanent_checkpoint``.
      eval_at: Function (integer --> boolean) that says, for training step n,
          whether that step should run evals. If ``None``, run evals on the
          first step and on every N'th step, as determined by the first
          training task.
      which_task: Function (integer --> integer) indicating which task should be
          used at which training step. Can be set to ``None`` in single-task
          training.
      n_devices: integer or ``None``, the number of devices for this
          computation.
      random_seed: the random seed to use; time/os dependent if ``None``
          (default).
      loss_chunk_size: int, if > 0 use chunks of this size to make loss
        computation more more memory-efficient.
      use_memory_efficient_trainer: whether to use a special memory-efficient
        trainer; if set to 2, the memory efficiency if very aggressive
      adasum: if True, use adaptive summation for multi-device gradients
      callbacks: List of subclasses of StepCallback to call on training
        steps.
    """
    self._is_chief, self._n_hosts, self._n_devices, self._rng = (
        init_host_and_devices(n_devices, random_seed))
    if use_memory_efficient_trainer:
      self._rng = tl.on_cpu(self._rng)

    # Handle single task case without lists too.
    if not isinstance(tasks, (list, tuple)):
      tasks = [tasks]

    if not tasks:
      raise ValueError('Must provide at least one training task.')
    if eval_tasks is None:
      eval_tasks = []
      eval_at = _never
    else:
      if not isinstance(eval_tasks, (list, tuple)):
        eval_tasks = [eval_tasks]

    self._tasks = tasks
    self._model = model
    self._eval_model = eval_model or model

    self._use_memory_efficient_trainer = use_memory_efficient_trainer
    self._loss_chunk_size = loss_chunk_size
    self._adasum = adasum
    # TODO(lukaszkaiser): can we have different eval models and save memory?
    if use_memory_efficient_trainer:
      assert len(tasks) == 1, 'only single task supported for now'
      self._eval_model = model

    default_at = _at_step_1_and_every_nth_step(tasks[0].n_steps_per_checkpoint)
    permanent_default_at = _at_step_1_and_every_nth_step(
        tasks[0].n_steps_per_permanent_checkpoint)
    if output_dir is not None:
      self._output_dir = os.path.expanduser(output_dir)
      tf.io.gfile.makedirs(self._output_dir)
      inputs.load_data_counters(self._output_dir)
    else:
      self._output_dir = None

    # Prepare training components.
    self._step = 0
    self._history = trax_history.History()
    self._checkpoint_at = checkpoint_at or default_at
    self._checkpoint_low_metric = checkpoint_low_metric
    self._checkpoint_high_metric = checkpoint_high_metric
    self._permanent_checkpoint_at = (
        permanent_checkpoint_at or permanent_default_at)
    if which_task is None:
      # If which task is not passed, then we permute tasks one by one.
      # If len(tasks) = 1, then which_task is a constant function equal to 0.
      which_task = lambda n: n % len(tasks)
    self._which_task = which_task

    # Initialize using the given random seed.
    # NOTE: If random_seed is None then self._rng will be different on
    # different hosts, leading to different weights on the different hosts.
    self._batch_signature = shapes.signature(tasks[0].sample_batch)
    self._model.rng = self.new_rng()
    # In the memory-efficient case, we initialize in init_trainer.
    if not use_memory_efficient_trainer:
      if _is_uninitialized(self._model):
        self._model.init(self._batch_signature)
      self._eval_model.rng = self.new_rng()
      if _is_uninitialized(self._eval_model):
        self._eval_model.init(self._batch_signature)

    # To handle the above case (i.e. random_seed = None), we psum the weights
    # and state and average them.
    # NOTE: This adds time (how much?) so we prefer not to do it if it is
    # unnecessary, i.e. random_seed was set.
    # NOTE: Averaging the weights across devices can screw up the initial weight
    # statistics.
    # TODO(pkozakowski): Broadcast from one of the devices instead?
    if (random_seed is None and self._n_hosts > 1 and
        not use_memory_efficient_trainer):
      logging.info('Syncing weights/state across %d hosts.', self._n_hosts)
      # Do self._sync_weights_and_state_across_hosts() but layer-by-layer
      # to save memory.
      blocks, last_layer = optimizers.trainer.extract_reversible_blocks(
          [self._model])
      all_layers = []
      for (std_layer, rev_layers) in blocks:
        all_layers.append(tl.Serial(std_layer))
        all_layers.extend(rev_layers)
      all_layers.append(last_layer)
      for layer in all_layers:
        weights_and_state = (layer.weights, layer.state)
        if not _is_empty(weights_and_state):
          layer.weights, layer.state = tl.on_cpu(self._unreplicate(
              _make_weights_and_state_same_across_hosts(
                  self._for_n_devices(weights_and_state))))

    # Create the optimizer for the training loss function.
    self._trainer_per_task = tuple(self._init_trainer(task) for task in tasks)

    # Sync layers weights/state in memory effcient trainer layers.
    if (random_seed is None and self._n_hosts > 1 and
        use_memory_efficient_trainer):
      logging.info('Syncing layers across %d hosts.', self._n_hosts)
      for layer in self._trainer_per_task[0].all_layers:
        weights_and_state = (layer.weights, layer.state)
        if not _is_empty(weights_and_state):
          layer.weights, layer.state = tl.on_cpu(self._unreplicate(
              _make_weights_and_state_same_across_hosts(
                  self._for_n_devices(weights_and_state))))

    # Load checkpoint if it exists.
    self.load_checkpoint()

    # Prepare eval components.
    self._eval_at = eval_at or default_at
    self._eval_tasks = eval_tasks
    loss_names = [task.loss_name for task in self._tasks]
    metric_names = [
        name  # pylint: disable=g-complex-comprehension
        for eval_task in self._eval_tasks
        for name in eval_task.metric_names
    ]
    self._rjust_len = max(map(len, loss_names + metric_names))
    self._evaluator_per_task = tuple(
        self._init_evaluator(eval_task) for eval_task in self._eval_tasks)

    if self._output_dir is None:
      _log('Will not write evaluation metrics, because output_dir is None.')

    def task_output_dir(task_index, task_list):
      if self._output_dir is not None:
        if len(task_list) < 2:
          output_dir = self._output_dir
        else:
          output_dir = os.path.join(
              self._output_dir,
              task_list[task_index].export_prefix or str(task_index))
        tf.io.gfile.makedirs(output_dir)
        return output_dir
      else:
        return None
    self._output_dir_per_eval_task = [
        task_output_dir(i, eval_tasks) for i in range(len(eval_tasks))]
    self._output_dir_per_train_task = [
        task_output_dir(i, tasks) for i in range(len(tasks))]

    callbacks = callbacks or []
    self._callbacks = [
        callback_class(self) for callback_class in callbacks
    ]
def __init__ ( self , model , tasks , eval_model = None , eval_tasks = None , output_dir = None , checkpoint_at = None , checkpoint_low_metric = None , checkpoint_high_metric = None , permanent_checkpoint_at = None , eval_at = None , which_task = None , n_devices = None , random_seed = None , loss_chunk_size = 0 , use_memory_efficient_trainer = False , adasum = False , callbacks = None , ) : self . _is_chief , self . _n_hosts , self . _n_devices , self . _rng = ( init_host_and_devices ( n_devices , random_seed ) ) if use_memory_efficient_trainer : self . _rng = tl . on_cpu ( self . _rng ) # Handle single task case without lists too. if not isinstance ( tasks , ( list , tuple ) ) : tasks = [ tasks ] if not tasks : raise ValueError ( 'Must provide at least one training task.' ) if eval_tasks is None : eval_tasks = [ ] eval_at = _never else : if not isinstance ( eval_tasks , ( list , tuple ) ) : eval_tasks = [ eval_tasks ] self . _tasks = tasks self . _model = model self . _eval_model = eval_model or model self . _use_memory_efficient_trainer = use_memory_efficient_trainer self . _loss_chunk_size = loss_chunk_size self . _adasum = adasum # TODO(lukaszkaiser): can we have different eval models and save memory? if use_memory_efficient_trainer : assert len ( tasks ) == 1 , 'only single task supported for now' self . _eval_model = model default_at = _at_step_1_and_every_nth_step ( tasks [ 0 ] . n_steps_per_checkpoint ) permanent_default_at = _at_step_1_and_every_nth_step ( tasks [ 0 ] . n_steps_per_permanent_checkpoint ) if output_dir is not None : self . _output_dir = os . path . expanduser ( output_dir ) tf . io . gfile . makedirs ( self . _output_dir ) inputs . load_data_counters ( self . _output_dir ) else : self . _output_dir = None # Prepare training components. self . _step = 0 self . _history = trax_history . History ( ) self . _checkpoint_at = checkpoint_at or default_at self . _checkpoint_low_metric = checkpoint_low_metric self . _checkpoint_high_metric = checkpoint_high_metric self . _permanent_checkpoint_at = ( permanent_checkpoint_at or permanent_default_at ) if which_task is None : # If which task is not passed, then we permute tasks one by one. # If len(tasks) = 1, then which_task is a constant function equal to 0. which_task = lambda n : n % len ( tasks ) self . _which_task = which_task # Initialize using the given random seed. # NOTE: If random_seed is None then self._rng will be different on # different hosts, leading to different weights on the different hosts. self . _batch_signature = shapes . signature ( tasks [ 0 ] . sample_batch ) self . _model . rng = self . new_rng ( ) # In the memory-efficient case, we initialize in init_trainer. if not use_memory_efficient_trainer : if _is_uninitialized ( self . _model ) : self . _model . init ( self . _batch_signature ) self . _eval_model . rng = self . new_rng ( ) if _is_uninitialized ( self . _eval_model ) : self . _eval_model . init ( self . _batch_signature ) # To handle the above case (i.e. random_seed = None), we psum the weights # and state and average them. # NOTE: This adds time (how much?) so we prefer not to do it if it is # unnecessary, i.e. random_seed was set. # NOTE: Averaging the weights across devices can screw up the initial weight # statistics. # TODO(pkozakowski): Broadcast from one of the devices instead? if ( random_seed is None and self . _n_hosts > 1 and not use_memory_efficient_trainer ) : logging . info ( 'Syncing weights/state across %d hosts.' , self . _n_hosts ) # Do self._sync_weights_and_state_across_hosts() but layer-by-layer # to save memory. blocks , last_layer = optimizers . trainer . extract_reversible_blocks ( [ self . _model ] ) all_layers = [ ] for ( std_layer , rev_layers ) in blocks : all_layers . append ( tl . Serial ( std_layer ) ) all_layers . extend ( rev_layers ) all_layers . append ( last_layer ) for layer in all_layers : weights_and_state = ( layer . weights , layer . state ) if not _is_empty ( weights_and_state ) : layer . weights , layer . state = tl . on_cpu ( self . _unreplicate ( _make_weights_and_state_same_across_hosts ( self . _for_n_devices ( weights_and_state ) ) ) ) # Create the optimizer for the training loss function. self . _trainer_per_task = tuple ( self . _init_trainer ( task ) for task in tasks ) # Sync layers weights/state in memory effcient trainer layers. if ( random_seed is None and self . _n_hosts > 1 and use_memory_efficient_trainer ) : logging . info ( 'Syncing layers across %d hosts.' , self . _n_hosts ) for layer in self . _trainer_per_task [ 0 ] . all_layers : weights_and_state = ( layer . weights , layer . state ) if not _is_empty ( weights_and_state ) : layer . weights , layer . state = tl . on_cpu ( self . _unreplicate ( _make_weights_and_state_same_across_hosts ( self . _for_n_devices ( weights_and_state ) ) ) ) # Load checkpoint if it exists. self . load_checkpoint ( ) # Prepare eval components. self . _eval_at = eval_at or default_at self . _eval_tasks = eval_tasks loss_names = [ task . loss_name for task in self . _tasks ] metric_names = [ name # pylint: disable=g-complex-comprehension for eval_task in self . _eval_tasks for name in eval_task . metric_names ] self . _rjust_len = max ( map ( len , loss_names + metric_names ) ) self . _evaluator_per_task = tuple ( self . _init_evaluator ( eval_task ) for eval_task in self . _eval_tasks ) if self . _output_dir is None : _log ( 'Will not write evaluation metrics, because output_dir is None.' ) def task_output_dir ( task_index , task_list ) : if self . _output_dir is not None : if len ( task_list ) < 2 : output_dir = self . _output_dir else : output_dir = os . path . join ( self . _output_dir , task_list [ task_index ] . export_prefix or str ( task_index ) ) tf . io . gfile . makedirs ( output_dir ) return output_dir else : return None self . _output_dir_per_eval_task = [ task_output_dir ( i , eval_tasks ) for i in range ( len ( eval_tasks ) ) ] self . _output_dir_per_train_task = [ task_output_dir ( i , tasks ) for i in range ( len ( tasks ) ) ] callbacks = callbacks or [ ] self . _callbacks = [ callback_class ( self ) for callback_class in callbacks ]