(self, n_steps=1)
Runs this training loop for n steps.

    Optionally runs evals and saves checkpoints at specified points.

    Args:
      n_steps: Stop training after completing n steps.
Runs this training loop for n steps.
Runs this training loop for n steps .
def run(self, n_steps=1):
    """Runs this training loop for n steps.

    Optionally runs evals and saves checkpoints at specified points.

    Args:
      n_steps: Stop training after completing n steps.
    """
    with self._open_summary_writers() as (
        train_summary_writers, eval_summary_writers):
      process = psutil.Process(os.getpid())
      loss_acc, step_acc = 0.0, 0
      start_time = time.time()
      optimizer_metrics_acc = collections.defaultdict(float)
      for i in range(n_steps):
        prev_task_index = self._which_task(self._step)
        self._step += 1
        task_index = self._which_task(self._step)
        task_changed = task_index != prev_task_index

        if task_changed:
          loss_acc, step_acc = 0.0, 0

        loss, optimizer_metrics = self._run_one_step(task_index, task_changed)

        # optimizer_metrics and loss are replicated on self.n_devices, a few
        # metrics are replicated (ex: gradients_l2, weights_l2) - i.e. they are
        # the same across devices, whereas some (ex: loss) aren't because they
        # are different on different devices (due to different data).
        # Taking the average does the correct thing in both the cases.
        #
        # NOTE: Only the weights and gradients are synced across the hosts. This
        # implies the loss here is averaged from this hosts' devices and not
        # across all hosts.
        optimizer_metrics, loss = fastmath.nested_map(
            functools.partial(tl.mean_or_pmean, self._n_devices),
            (optimizer_metrics, loss))

        loss_acc += loss
        # Log loss every 50 steps, every step in memory-efficient trainer.
        if self._step % 50 == 0 or self._use_memory_efficient_trainer:
          self._log_step('Loss: %.4f' % loss, stdout=False)
        step_acc += 1
        for metric_name, value in optimizer_metrics.items():
          optimizer_metrics_acc[metric_name] += value

        # TODO(yuwenyan): Finds a way to log the last round eval step in
        # history.
        #
        # Right now, the last round eval log is missing in history since the
        # checkpoint is saved before it. However sometimes the eval step will
        # fail for some reasons, and it's not acceptable to loose the whole
        # checkpoint in this case. Stays with the old way for now, and fixes it
        # when the checkpoint format is changed to storing weights separately
        # from a small file with history and other data.
        if self._checkpoint_at(self.step):
          self.save_checkpoint('model')
        if self._permanent_checkpoint_at(self.step):
          self.save_checkpoint(f'model_{self.step}')
        if self._eval_at(self.step):
          logging.info('cpu memory use (MB): %.2f',
                       process.memory_info().rss / float(1024 * 1024))
          elapsed_time = time.time() - start_time
          self._log_training_progress(
              task=self._tasks[task_index],
              total_loss=loss_acc,
              n_steps=step_acc,
              elapsed_time=elapsed_time,
              optimizer_metrics=optimizer_metrics_acc,
              summary_writer=train_summary_writers[task_index],
          )
          self.run_evals(eval_summary_writers)
          loss_acc, step_acc = 0.0, 0
          start_time = time.time()
          optimizer_metrics_acc = collections.defaultdict(float)

        # For the current step, after all evals are run and recorded in the
        # event history, check if we need to save special checkpoints because
        # of a new low metric value or a new high metric value.
        if self._checkpoint_at(self.step):
          if self._checkpoint_low_metric is not None and self._at_lowest():
            self.save_checkpoint(f'lowest_{self._checkpoint_low_metric}')
          if self._checkpoint_high_metric is not None and self._at_highest():
            self.save_checkpoint(f'highest_{self._checkpoint_high_metric}')

    # Store the final values back into their respective objects, for testing
    # or other inspection/use.
    #
    # We keep the standard model weights/state unreplicated and
    # tl.Accelerate(model) will carry the replicated weights/state.
    # TODO(afrozm): Try to use tl.Accelerate(model) everywhere in the Loop.
    self._eval_model.weights = self._model.weights
def run ( self , n_steps = 1 ) : with self . _open_summary_writers ( ) as ( train_summary_writers , eval_summary_writers ) : process = psutil . Process ( os . getpid ( ) ) loss_acc , step_acc = 0.0 , 0 start_time = time . time ( ) optimizer_metrics_acc = collections . defaultdict ( float ) for i in range ( n_steps ) : prev_task_index = self . _which_task ( self . _step ) self . _step += 1 task_index = self . _which_task ( self . _step ) task_changed = task_index != prev_task_index if task_changed : loss_acc , step_acc = 0.0 , 0 loss , optimizer_metrics = self . _run_one_step ( task_index , task_changed ) # optimizer_metrics and loss are replicated on self.n_devices, a few # metrics are replicated (ex: gradients_l2, weights_l2) - i.e. they are # the same across devices, whereas some (ex: loss) aren't because they # are different on different devices (due to different data). # Taking the average does the correct thing in both the cases. # # NOTE: Only the weights and gradients are synced across the hosts. This # implies the loss here is averaged from this hosts' devices and not # across all hosts. optimizer_metrics , loss = fastmath . nested_map ( functools . partial ( tl . mean_or_pmean , self . _n_devices ) , ( optimizer_metrics , loss ) ) loss_acc += loss # Log loss every 50 steps, every step in memory-efficient trainer. if self . _step % 50 == 0 or self . _use_memory_efficient_trainer : self . _log_step ( 'Loss: %.4f' % loss , stdout = False ) step_acc += 1 for metric_name , value in optimizer_metrics . items ( ) : optimizer_metrics_acc [ metric_name ] += value # TODO(yuwenyan): Finds a way to log the last round eval step in # history. # # Right now, the last round eval log is missing in history since the # checkpoint is saved before it. However sometimes the eval step will # fail for some reasons, and it's not acceptable to loose the whole # checkpoint in this case. Stays with the old way for now, and fixes it # when the checkpoint format is changed to storing weights separately # from a small file with history and other data. if self . _checkpoint_at ( self . step ) : self . save_checkpoint ( 'model' ) if self . _permanent_checkpoint_at ( self . step ) : self . save_checkpoint ( f'model_{self.step}' ) if self . _eval_at ( self . step ) : logging . info ( 'cpu memory use (MB): %.2f' , process . memory_info ( ) . rss / float ( 1024 * 1024 ) ) elapsed_time = time . time ( ) - start_time self . _log_training_progress ( task = self . _tasks [ task_index ] , total_loss = loss_acc , n_steps = step_acc , elapsed_time = elapsed_time , optimizer_metrics = optimizer_metrics_acc , summary_writer = train_summary_writers [ task_index ] , ) self . run_evals ( eval_summary_writers ) loss_acc , step_acc = 0.0 , 0 start_time = time . time ( ) optimizer_metrics_acc = collections . defaultdict ( float ) # For the current step, after all evals are run and recorded in the # event history, check if we need to save special checkpoints because # of a new low metric value or a new high metric value. if self . _checkpoint_at ( self . step ) : if self . _checkpoint_low_metric is not None and self . _at_lowest ( ) : self . save_checkpoint ( f'lowest_{self._checkpoint_low_metric}' ) if self . _checkpoint_high_metric is not None and self . _at_highest ( ) : self . save_checkpoint ( f'highest_{self._checkpoint_high_metric}' ) # Store the final values back into their respective objects, for testing # or other inspection/use. # # We keep the standard model weights/state unreplicated and # tl.Accelerate(model) will carry the replicated weights/state. # TODO(afrozm): Try to use tl.Accelerate(model) everywhere in the Loop. self . _eval_model . weights = self . _model . weights