Store variable scopes in T2TModel; add T2TModel.initialize_from_ckpt

PiperOrigin-RevId: 209218783

Author: Ryan Sepassi

tensor2tensor/bin/t2t_trainer.py

@@ -180,6 +180,7 @@ def create_experiment_fn(**kwargs):
       use_tpu=FLAGS.use_tpu,
       use_tpu_estimator=FLAGS.use_tpu_estimator,
       use_xla=FLAGS.xla_compile,
+      warm_start_from=FLAGS.warm_start_from,
       **kwargs)
 
 
@@ -214,7 +215,6 @@ def create_run_config(hp):
       hp.weight_dtype == "float32")
   return trainer_lib.create_run_config(
       model_dir=os.path.expanduser(FLAGS.output_dir),
-      warm_start_from=FLAGS.warm_start_from,
       master=FLAGS.master,
       iterations_per_loop=FLAGS.iterations_per_loop,
       num_shards=FLAGS.tpu_num_shards,

tensor2tensor/utils/t2t_model.py

@@ -120,6 +120,11 @@ def __init__(self,
       self._create_modalities(self._problem_hparams, self._hparams)
     if not common_layers.is_xla_compiled():
       self.summarize_hparams()
+    self._variable_scopes = {}
+
+  def _add_variable_scope(self, key, vs):
+    if key not in self._variable_scopes:
+      self._variable_scopes[key] = vs
 
   def summarize_hparams(self):
     def create_hparams_summary(hparams, name):
@@ -261,15 +266,17 @@ def model_fn_sharded(self, sharded_features):
     return sharded_logits, losses
 
   def model_fn(self, features):
-    with tf.variable_scope(tf.get_variable_scope(), use_resource=True):
+    with tf.variable_scope(tf.get_variable_scope(), use_resource=True) as vs:
+      self._add_variable_scope("model_fn", vs)
       transformed_features = self.bottom(features)
 
       if self.hparams.activation_dtype == "bfloat16":
         for k, v in sorted(six.iteritems(transformed_features)):
           if v.dtype == tf.float32:
             transformed_features[k] = tf.cast(v, tf.bfloat16)
 
-      with tf.variable_scope("body"):
+      with tf.variable_scope("body") as body_vs:
+        self._add_variable_scope("body", body_vs)
         log_info("Building model body")
         body_out = self.body(transformed_features)
       output, losses = self._normalize_body_output(body_out)
@@ -302,7 +309,8 @@ def bottom(self, features):
         tf.logging.warning("Missing feature %s - ignoring." % key)
         continue
       do_reuse = input_modality.name in all_previous_modalities
-      with tf.variable_scope(input_modality.name, reuse=do_reuse):
+      with tf.variable_scope(input_modality.name, reuse=do_reuse) as im_vs:
+        self._add_variable_scope(input_modality.name, im_vs)
         log_info("Transforming feature '%s' with %s.bottom", key,
                  input_modality.name)
         transformed_features[key] = input_modality.bottom(features[key])
@@ -313,14 +321,16 @@ def bottom(self, features):
     if isinstance(target_modality, dict):
       for k, v in six.iteritems(target_modality):
         if k in features:
-          with tf.variable_scope(
-              "%s/%s" % (v.name, k)):  # TODO(aidangomez): share variables?
+          # TODO(aidangomez): share variables?
+          with tf.variable_scope("%s/%s" % (v.name, k)) as tm_vs:
+            self._add_variable_scope("%s/%s" % (v.name, k), tm_vs)
             log_info("Transforming '%s' with %s.targets_bottom", k, v.name)
             transformed_features[k] = v.targets_bottom(features[k])
         else:
           tf.logging.warn("Modality not found in features: %s", k)
     else:
-      with tf.variable_scope(target_modality.name):
+      with tf.variable_scope(target_modality.name) as tm_vs:
+        self._add_variable_scope(target_modality.name, tm_vs)
         if "targets" in features:
           log_info("Transforming 'targets' with %s.targets_bottom",
                    target_modality.name)
@@ -359,7 +369,8 @@ def _top_single(self, body_output, target_modality, features):
       log_warn("Without a Problem, T2TModel.top is a passthrough.")
       return body_output
 
-    with tf.variable_scope(target_modality.name):
+    with tf.variable_scope(target_modality.name) as tm_vs:
+      self._add_variable_scope(tm_vs.name, tm_vs)
       log_info("Transforming body output with %s.top", target_modality.name)
       last_only = (
           target_modality.top_is_pointwise and
@@ -401,7 +412,9 @@ def top(self, body_output, features):
             "problem_hparams.target_modality's dict." % k)
       logits = {}
       for k, v in six.iteritems(body_output):
-        with tf.variable_scope(k):  # TODO(aidangomez): share variables here?
+        # TODO(aidangomez): share variables here?
+        with tf.variable_scope(k) as top_vs:
+          self._add_variable_scope("top_%s" % k, top_vs)
           logits[k] = self._top_single(v, target_modality[k], features)
       return logits
     else:
@@ -1270,26 +1283,33 @@ def estimator_model_fn(cls,
     return model.estimator_spec_train(
         loss, num_async_replicas=num_async_replicas, use_tpu=use_tpu)
 
+  def initialize_from_ckpt(self, ckpt_dir):
+    model_dir = self._hparams.get("model_dir", None)
+    already_has_ckpt = (
+        model_dir and tf.train.latest_checkpoint(model_dir) is not None)
+    if already_has_ckpt:
+      return
+
+    # TODO(mitchellstern): Add support for partitioned variables?
+    reader = tf.contrib.framework.load_checkpoint(ckpt_dir)
+    variable_map = {}
+    for var in tf.contrib.framework.get_trainable_variables():
+      var_name = var.name.split(":")[0]
+      if reader.has_tensor(var_name):
+        tf.logging.info("Loading variable from checkpoint: %s", var_name)
+        variable_map[var_name] = var
+      else:
+        tf.logging.info(
+            "Cannot find variable in checkpoint, skipping: %s", var_name)
+    tf.train.init_from_checkpoint(ckpt_dir, variable_map)
+
   def estimator_spec_train(self, loss, num_async_replicas=1, use_tpu=False):
     """Construct EstimatorSpec for TRAIN mode."""
     train_op = self.optimize(loss, num_async_replicas=num_async_replicas,
                              use_tpu=use_tpu)
 
-    # TODO(mitchellstern): Add support for partitioned variables?
-    if (tf.train.latest_checkpoint(self._hparams.model_dir) is None and
-        self._hparams.pretrained_model_dir):
-      pretrained_model_dir = self._hparams.pretrained_model_dir
-      reader = tf.contrib.framework.load_checkpoint(pretrained_model_dir)
-      variable_map = {}
-      for var in tf.contrib.framework.get_trainable_variables():
-        var_name = var.name.split(":")[0]
-        if reader.has_tensor(var_name):
-          tf.logging.info("Loading variable from checkpoint: %s", var_name)
-          variable_map[var_name] = var
-        else:
-          tf.logging.info(
-              "Cannot find variable in checkpoint, skipping: %s", var_name)
-      tf.train.init_from_checkpoint(pretrained_model_dir, variable_map)
+    if self._hparams.warm_start_from:
+      self.initialize_from_ckpt(self._hparams.warm_start_from)
 
     if use_tpu:
       host_call = _create_host_call(self.hparams.model_dir)

tensor2tensor/utils/trainer_lib.py

@@ -105,7 +105,6 @@ def is_cloud_async_distributed():
 
 def create_run_config(master="",
                       model_dir=None,
-                      warm_start_from=None,
                       iterations_per_loop=1000,
                       num_shards=8,
                       log_device_placement=False,
@@ -197,7 +196,6 @@ def create_run_config(master="",
     del run_config_args["evaluation_master"]
 
   config = run_config_cls(**run_config_args)
-  config.warm_start_from = warm_start_from
 
   # If not using TPU, add device info for data_parallelism
   config.use_tpu = use_tpu
@@ -259,7 +257,6 @@ def create_estimator(model_name,
         model_fn=model_fn,
         model_dir=run_config.model_dir,
         config=run_config,
-        warm_start_from=run_config.warm_start_from
     )
   return estimator
 
@@ -432,14 +429,16 @@ def create_experiment(
     use_tpu_estimator=False,
     use_xla=False,
     additional_train_hooks=None,
-    additional_eval_hooks=None):
+    additional_eval_hooks=None,
+    warm_start_from=None):
   """Create Experiment."""
   # HParams
   hparams.add_hparam("model_dir", run_config.model_dir)
   hparams.add_hparam("data_dir", data_dir)
   hparams.add_hparam("train_steps", train_steps)
   hparams.add_hparam("eval_steps", eval_steps)
   hparams.add_hparam("schedule", schedule)
+  hparams.add_hparam("warm_start_from", warm_start_from)
   add_problem_hparams(hparams, problem_name)
 
   # Estimator