Method test_recovery

official/core/train_lib_test.py:222–260

(self, distribution_strategy, flag_mode)

          flag_mode=['train'],
      ))
  def test_recovery(self, distribution_strategy, flag_mode):
    loss_threshold = 1.0
    model_dir = self.get_temp_dir()
    flags_dict = dict(
        experiment='mock',
        mode=flag_mode,
        model_dir=model_dir,
        params_override=json.dumps(self._test_config))
    with flagsaver.flagsaver(**flags_dict):
      params = train_utils.parse_configuration(flags.FLAGS)
      params.trainer.loss_upper_bound = loss_threshold
      params.trainer.recovery_max_trials = 1
      train_utils.serialize_config(params, model_dir)
      with distribution_strategy.scope():
        task = task_factory.get_task(params.task, logging_dir=model_dir)

      # Saves a checkpoint for reference.
      model = task.build_model()
      checkpoint = tf.train.Checkpoint(model=model)
      checkpoint_manager = tf.train.CheckpointManager(
          checkpoint, self.get_temp_dir(), max_to_keep=2)
      checkpoint_manager.save()
      before_weights = model.get_weights()

      def build_losses(labels, model_outputs, aux_losses=None):
        del labels, model_outputs
        return tf.constant([loss_threshold], tf.float32) + aux_losses

      task.build_losses = build_losses

      model, _ = train_lib.OrbitExperimentRunner(
          distribution_strategy=distribution_strategy,
          task=task,
          mode=flag_mode,
          params=params,
          model_dir=model_dir).run()
      after_weights = model.get_weights()
      for left, right in zip(before_weights, after_weights):
        self.assertAllEqual(left, right)

  def test_parse_configuration(self):
    model_dir = self.get_temp_dir()