Python tensor2tensor.utils.registry.model() Examples

def __init__(self,
               hparams,
               mode=tf.estimator.ModeKeys.TRAIN,
               problem_hparams=None,
               data_parallelism=None,
               decode_hparams=None,
               **kwargs):
    assert hparams.distill_phase in ["train", "distill"]

    if hparams.distill_phase == "train" and hparams.teacher_learning_rate:
      hparams.learning_rate = hparams.teacher_learning_rate
    elif hparams.distill_phase == "distill" and hparams.student_learning_rate:
      hparams.learning_rate = hparams.student_learning_rate

    self.teacher_hparams = registry.hparams(hparams.teacher_hparams)
    self.teacher_model = registry.model(
        hparams.teacher_model)(self.teacher_hparams, mode, problem_hparams,
                               data_parallelism, decode_hparams)
    self.student_hparams = registry.hparams(hparams.student_hparams)
    self.student_model = registry.model(
        hparams.student_model)(self.student_hparams, mode, problem_hparams,
                               data_parallelism, decode_hparams)
    super(Distillation,
          self).__init__(hparams, mode, problem_hparams, data_parallelism,
                         decode_hparams, **kwargs) 

def evaluate_world_model(simulated_problem_name, problem_name, hparams,
                         world_model_dir, epoch_data_dir, tmp_dir):
  """Generate simulated environment data and return reward accuracy."""
  gym_simulated_problem = registry.problem(simulated_problem_name)
  sim_steps = hparams.simulated_env_generator_num_steps
  gym_simulated_problem.settable_num_steps = sim_steps
  with temporary_flags({
      "problem": problem_name,
      "model": hparams.generative_model,
      "hparams_set": hparams.generative_model_params,
      "data_dir": epoch_data_dir,
      "output_dir": world_model_dir,
  }):
    gym_simulated_problem.generate_data(epoch_data_dir, tmp_dir)
  n = max(1., gym_simulated_problem.statistics.number_of_dones)
  model_reward_accuracy = (
      gym_simulated_problem.statistics.successful_episode_reward_predictions
      / float(n))
  old_path = os.path.join(epoch_data_dir, "debug_frames_sim")
  new_path = os.path.join(epoch_data_dir, "debug_frames_sim_eval")
  if not tf.gfile.Exists(new_path):
    tf.gfile.Rename(old_path, new_path)
  return model_reward_accuracy 

def make_estimator_model_fn(model_name,
                              hparams,
                              decode_hparams=None,
                              use_tpu=False):
    model_cls = registry.model(model_name)

    def wrapping_model_fn(features, labels, mode, params=None, config=None):
      return model_cls.estimator_model_fn(
          hparams,
          features,
          labels,
          mode,
          config=config,
          params=params,
          decode_hparams=decode_hparams,
          use_tpu=use_tpu)

    return wrapping_model_fn 

def __init__(self,
               hparams,
               mode=tf.estimator.ModeKeys.TRAIN,
               problem_hparams=None,
               data_parallelism=None,
               decode_hparams=None):
    assert hparams.distill_phase in ["train", "distill"]

    if hparams.distill_phase == "train" and hparams.teacher_learning_rate:
      hparams.learning_rate = hparams.teacher_learning_rate
    elif hparams.distill_phase == "distill" and hparams.student_learning_rate:
      hparams.learning_rate = hparams.student_learning_rate

    self.teacher_hparams = registry.hparams(hparams.teacher_hparams)
    self.teacher_model = registry.model(
        hparams.teacher_model)(self.teacher_hparams, mode, problem_hparams,
                               data_parallelism, decode_hparams)
    self.student_hparams = registry.hparams(hparams.student_hparams)
    self.student_model = registry.model(
        hparams.student_model)(self.student_hparams, mode, problem_hparams,
                               data_parallelism, decode_hparams)
    super(Distillation, self).__init__(hparams, mode, problem_hparams,
                                       data_parallelism, decode_hparams) 

def make_estimator_model_fn(model_name,
                              hparams,
                              decode_hparams=None,
                              use_tpu=False):
    model_cls = registry.model(model_name)

    def wrapping_model_fn(features, labels, mode, params=None, config=None):
      return model_cls.estimator_model_fn(
          hparams,
          features,
          labels,
          mode,
          config=config,
          params=params,
          decode_hparams=decode_hparams,
          use_tpu=use_tpu)

    return wrapping_model_fn 

def ppo_original_params():
  """Parameters based on the original PPO paper."""
  hparams = ppo_atari_base()
  hparams.learning_rate_constant = 2.5e-4
  hparams.gae_gamma = 0.99
  hparams.gae_lambda = 0.95
  hparams.clipping_coef = 0.1
  hparams.value_loss_coef = 1
  hparams.entropy_loss_coef = 0.01
  hparams.eval_every_epochs = 200
  hparams.dropout_ppo = 0.1
  # The parameters below are modified to accommodate short epoch_length (which
  # is needed for model based rollouts).
  hparams.epoch_length = 50
  hparams.optimization_batch_size = 20
  return hparams 

def ppo_original_params():
  """Parameters based on the original PPO paper."""
  hparams = ppo_atari_base()
  hparams.learning_rate_constant = 2.5e-4
  hparams.gae_gamma = 0.99
  hparams.gae_lambda = 0.95
  hparams.clipping_coef = 0.1
  hparams.value_loss_coef = 1
  hparams.entropy_loss_coef = 0.01
  hparams.eval_every_epochs = 200
  hparams.dropout_ppo = 0.1
  # The parameters below are modified to accommodate short epoch_length (which
  # is needed for model based rollouts).
  hparams.epoch_length = 50
  hparams.optimization_batch_size = 20
  return hparams 

def get_mnist_random_output(self, model_name, hparams_set=None,
                              mode=tf.estimator.ModeKeys.TRAIN):
    hparams_set = hparams_set or model_name
    x = np.random.randint(256, size=(1, 28, 28, 1))
    y = np.random.randint(10, size=(1, 1))
    features = {
        "targets": tf.constant(x, dtype=tf.int32),
        "inputs": tf.constant(y, dtype=tf.int32),
    }
    hparams = trainer_lib.create_hparams(
        hparams_set, problem_name="image_mnist_rev", data_dir=".")
    model = registry.model(model_name)(hparams, mode)
    tf.train.create_global_step()
    logits, _ = model(features)
    with self.test_session() as session:
      session.run(tf.global_variables_initializer())
      res = session.run(logits)
    return res 

def make_estimator_model_fn(model_name,
                              hparams,
                              decode_hparams=None,
                              use_tpu=False):
    model_cls = registry.model(model_name)

    def wrapping_model_fn(features, labels, mode, params=None, config=None):
      return model_cls.estimator_model_fn(
          hparams,
          features,
          labels,
          mode,
          config=config,
          params=params,
          decode_hparams=decode_hparams,
          use_tpu=use_tpu)

    return wrapping_model_fn 

def ppo_original_world_model():
  """Atari parameters with world model as policy."""
  hparams = ppo_original_params()
  hparams.policy_network = "next_frame_basic_deterministic"
  hparams_keys = hparams.values().keys()
  video_hparams = basic_deterministic_params.next_frame_basic_deterministic()
  for (name, value) in six.iteritems(video_hparams.values()):
    if name in hparams_keys:
      hparams.set_hparam(name, value)
    else:
      hparams.add_hparam(name, value)
  # Mostly to avoid decaying WM params when training the policy.
  hparams.weight_decay = 0
  return hparams 

def ppo_tiny_world_model():
  """Atari parameters with world model as policy."""
  hparams = ppo_original_params()
  hparams.policy_network = "next_frame_basic_deterministic"
  hparams_keys = hparams.values().keys()
  video_hparams = basic_deterministic_params.next_frame_tiny()
  for (name, value) in six.iteritems(video_hparams.values()):
    if name in hparams_keys:
      hparams.set_hparam(name, value)
    else:
      hparams.add_hparam(name, value)
  hparams.weight_decay = 0
  return hparams 

def estimator_spec_train(self, loss, num_async_replicas=1, use_tpu=False):
    """Constructs `tf.estimator.EstimatorSpec` for TRAIN (training) mode."""
    train_op = self.optimize(loss, num_async_replicas=num_async_replicas,
                             use_tpu=use_tpu)

    if use_tpu:
      if self._hparams.warm_start_from:
        def scaffold_fn():
          self.initialize_from_ckpt(self._hparams.warm_start_from)
          return tf.train.Scaffold()
      else:
        scaffold_fn = None

      # Note: important to call this before remove_summaries()
      if self.hparams.tpu_enable_host_call:
        host_call = self.create_train_host_call()
      else:
        host_call = None

      remove_summaries()

      return tf.contrib.tpu.TPUEstimatorSpec(
          tf.estimator.ModeKeys.TRAIN,
          loss=loss,
          train_op=train_op,
          host_call=host_call,
          scaffold_fn=scaffold_fn)
    else:
      if self._hparams.warm_start_from:
        self.initialize_from_ckpt(self._hparams.warm_start_from)

      # When loading weights from a pre-trained model, you want to be able to
      # load separate weights into the encoder and decoder.
      if self._hparams.warm_start_from_second:
        self.initialize_from_ckpt(self._hparams.warm_start_from_second)

      return tf.estimator.EstimatorSpec(
          tf.estimator.ModeKeys.TRAIN, loss=loss, train_op=train_op) 

def testModel(self):
    # HParams
    hparams = trainer_lib.create_hparams(
        "transformer_tiny", data_dir=algorithmic.TinyAlgo.data_dir,
        problem_name="tiny_algo")

    # Dataset
    problem = hparams.problem
    dataset = problem.dataset(tf.estimator.ModeKeys.TRAIN,
                              algorithmic.TinyAlgo.data_dir)
    dataset = dataset.repeat(None).padded_batch(10, dataset.output_shapes)
    features = dataset.make_one_shot_iterator().get_next()
    features = data_reader.standardize_shapes(features)

    # Model
    model = registry.model("transformer")(hparams, tf.estimator.ModeKeys.TRAIN)
    logits, losses = model(features)

    self.assertTrue("training" in losses)
    loss = losses["training"]

    with self.test_session() as sess:
      sess.run(tf.global_variables_initializer())
      logits_val, loss_val = sess.run([logits, loss])
      logits_shape = list(logits_val.shape)
      logits_shape[1] = None
      self.assertAllEqual(logits_shape, [10, None, 1, 1, 4])
      self.assertEqual(loss_val.shape, tuple()) 

def rlmf_tiny():
  """Tiny set of hparams for model-free PPO."""
  hparams = rlmf_original()
  hparams = hparams.override_from_dict(rlmf_tiny_overrides())
  hparams.batch_size = 2
  hparams.base_algo_params = "ppo_original_tiny"
  hparams.add_hparam("ppo_epochs_num", 3)
  hparams.add_hparam("ppo_epoch_length", 2)
  return hparams 

def rlmf_base():
  """Base set of hparams for model-free PPO."""
  hparams = rlmf_original()
  hparams.add_hparam("ppo_epochs_num", 3000)
  hparams.add_hparam("ppo_eval_every_epochs", 100)
  return hparams 

def train_autoencoder(problem_name, data_dir, output_dir, hparams, epoch):
  """Train autoencoder on problem_name."""
  train_steps = hparams.autoencoder_train_steps * (epoch + 2)
  with temporary_flags({
      "problem": problem_name,
      "data_dir": data_dir,
      "output_dir": output_dir,
      "model": "autoencoder_ordered_discrete",
      "hparams_set": "autoencoder_discrete_pong",
      "train_steps": train_steps,
      "eval_steps": 100,
  }):
    t2t_trainer.main([]) 

def get_eval_hooks(model_name, hook_context):
    model_cls = registry.model(model_name)
    return model_cls.eval_hooks(hook_context) 

def ppo_original_world_model_stochastic_discrete():
  """Atari parameters with stochastic discrete world model as policy."""
  hparams = ppo_original_params()
  hparams.policy_network = "next_frame_basic_stochastic_discrete"
  hparams_keys = hparams.values().keys()
  video_hparams = basic_stochastic.next_frame_basic_stochastic_discrete()
  for (name, value) in six.iteritems(video_hparams.values()):
    if name in hparams_keys:
      hparams.set_hparam(name, value)
    else:
      hparams.add_hparam(name, value)
  # To avoid OOM. Probably way to small.
  hparams.optimization_batch_size = 1
  hparams.weight_decay = 0
  return hparams 

def rlmf_eval_dist_threshold():
  """Distributional set of hparams for model-free PPO."""
  hparams = rlmf_eval_dist()
  hparams.distributional_threshold = 0.5
  return hparams 

def rlmf_eval():
  """Eval set of hparams for model-free PPO."""
  hparams = rlmf_original()
  hparams.batch_size = 16
  hparams.eval_batch_size = 32
  hparams.eval_episodes_num = 2
  hparams.eval_sampling_temps = [0.5, 0.0, 1.0]
  hparams.eval_rl_env_max_episode_steps = 40000
  hparams.add_hparam("ppo_epoch_length", 128)
  hparams.add_hparam("ppo_optimization_batch_size", 32)
  hparams.add_hparam("ppo_epochs_num", 10000)
  hparams.add_hparam("ppo_eval_every_epochs", 500)
  hparams.add_hparam("attempt", 0)
  hparams.add_hparam("moe_loss_coef", 0)
  return hparams 

def get_eval_hooks(model_name, hook_context):
    model_cls = registry.model(model_name)
    return model_cls.eval_hooks(hook_context) 

def rlmf_dist_threshold():
  """Distributional set of hparams for model-free PPO."""
  hparams = rlmf_dist()
  hparams.distributional_threshold = 0.5
  return hparams 

def rlmf_dist():
  """Distributional set of hparams for model-free PPO."""
  hparams = rlmf_original()
  hparams.distributional_size = 1024
  hparams.base_algo_params = "ppo_dist_params"
  return hparams 

def rlmf_base():
  """Base set of hparams for model-free PPO."""
  hparams = rlmf_original()
  hparams.add_hparam("ppo_epochs_num", 3000)
  hparams.add_hparam("ppo_eval_every_epochs", 100)
  return hparams 

def ppo_original_world_model_stochastic_discrete():
  """Atari parameters with stochastic discrete world model as policy."""
  hparams = ppo_original_params()
  hparams.policy_network = "next_frame_basic_stochastic_discrete"
  hparams_keys = hparams.values().keys()
  video_hparams = basic_stochastic.next_frame_basic_stochastic_discrete()
  for (name, value) in six.iteritems(video_hparams.values()):
    if name in hparams_keys:
      hparams.set_hparam(name, value)
    else:
      hparams.add_hparam(name, value)
  # To avoid OOM. Probably way to small.
  hparams.optimization_batch_size = 1
  hparams.weight_decay = 0
  return hparams 

def ppo_tiny_world_model():
  """Atari parameters with world model as policy."""
  hparams = ppo_original_params()
  hparams.policy_network = "next_frame_basic_deterministic"
  hparams_keys = hparams.values().keys()
  video_hparams = basic_deterministic_params.next_frame_tiny()
  for (name, value) in six.iteritems(video_hparams.values()):
    if name in hparams_keys:
      hparams.set_hparam(name, value)
    else:
      hparams.add_hparam(name, value)
  hparams.weight_decay = 0
  return hparams 

def ppo_original_world_model():
  """Atari parameters with world model as policy."""
  hparams = ppo_original_params()
  hparams.policy_network = "next_frame_basic_deterministic"
  hparams_keys = hparams.values().keys()
  video_hparams = basic_deterministic_params.next_frame_basic_deterministic()
  for (name, value) in six.iteritems(video_hparams.values()):
    if name in hparams_keys:
      hparams.set_hparam(name, value)
    else:
      hparams.add_hparam(name, value)
  # Mostly to avoid decaying WM params when training the policy.
  hparams.weight_decay = 0
  return hparams 

def testModel(self):
    # HParams
    hparams = trainer_lib.create_hparams(
        "transformer_tiny", data_dir=algorithmic.TinyAlgo.data_dir,
        problem_name="tiny_algo")

    # Dataset
    problem = hparams.problem
    dataset = problem.dataset(tf.estimator.ModeKeys.TRAIN,
                              algorithmic.TinyAlgo.data_dir)
    dataset = dataset.repeat(None).padded_batch(10, dataset.output_shapes)
    features = dataset.make_one_shot_iterator().get_next()
    features = data_reader.standardize_shapes(features)

    # Model
    model = registry.model("transformer")(hparams, tf.estimator.ModeKeys.TRAIN)
    logits, losses = model(features)

    self.assertTrue("training" in losses)
    loss = losses["training"]

    with self.test_session() as sess:
      sess.run(tf.global_variables_initializer())
      logits_val, loss_val = sess.run([logits, loss])
      logits_shape = list(logits_val.shape)
      logits_shape[1] = None
      self.assertAllEqual(logits_shape, [10, None, 1, 1, 4])
      self.assertEqual(loss_val.shape, tuple()) 

def estimator_spec_train(self, loss, num_async_replicas=1, use_tpu=False):
    """Constructs `tf.estimator.EstimatorSpec` for TRAIN (training) mode."""
    train_op = self.optimize(loss, num_async_replicas=num_async_replicas,
                             use_tpu=use_tpu)

    if use_tpu:
      if self._hparams.warm_start_from:
        def scaffold_fn():
          self.initialize_from_ckpt(self._hparams.warm_start_from)
          return tf.train.Scaffold()
      else:
        scaffold_fn = None

      # Note: important to call this before remove_summaries()
      if self.hparams.tpu_enable_host_call:
        host_call = self.create_train_host_call()
      else:
        host_call = None

      remove_summaries()

      return contrib.tpu().TPUEstimatorSpec(
          tf.estimator.ModeKeys.TRAIN,
          loss=loss,
          train_op=train_op,
          host_call=host_call,
          scaffold_fn=scaffold_fn)
    else:
      if self._hparams.warm_start_from:
        self.initialize_from_ckpt(self._hparams.warm_start_from)

      # When loading weights from a pre-trained model, you want to be able to
      # load separate weights into the encoder and decoder.
      if self._hparams.warm_start_from_second:
        self.initialize_from_ckpt(self._hparams.warm_start_from_second)

      return tf.estimator.EstimatorSpec(
          tf.estimator.ModeKeys.TRAIN, loss=loss, train_op=train_op) 

def rlmf_dist():
  """Distributional set of hparams for model-free PPO."""
  hparams = rlmf_original()
  hparams.distributional_size = 1024
  hparams.base_algo_params = "ppo_dist_params"
  return hparams