Merge pull request #694 from SkyAndCloud/raw

add support for specifying checkpoint_path in t2t-decoder interactive mode

Author: lukaszkaiser

tensor2tensor/bin/t2t_decoder.py

@@ -82,7 +82,7 @@ def create_decode_hparams():
 
 def decode(estimator, hparams, decode_hp):
   if FLAGS.decode_interactive:
-    decoding.decode_interactively(estimator, hparams, decode_hp)
+    decoding.decode_interactively(estimator, hparams, decode_hp, checkpoint_path=FLAGS.checkpoint_path)
   elif FLAGS.decode_from_file:
     decoding.decode_from_file(estimator, FLAGS.decode_from_file, hparams,
                               decode_hp, FLAGS.decode_to_file,

tensor2tensor/bin/t2t_trainer.py

@@ -178,7 +178,7 @@ def create_run_config(hp):
   save_ckpt_secs = FLAGS.save_checkpoints_secs or None
   if save_ckpt_secs:
     save_ckpt_steps = None
-  assert FLAGS.output_dir
+  assert FLAGS.output_dir or FLAGS.checkpoint_path
   return trainer_lib.create_run_config(
       model_dir=os.path.expanduser(FLAGS.output_dir),
       master=FLAGS.master,

tensor2tensor/utils/decoding.py

@@ -343,7 +343,7 @@ def input_fn():
   return input_fn
 
 
-def decode_interactively(estimator, hparams, decode_hp):
+def decode_interactively(estimator, hparams, decode_hp, checkpoint_path=None):
   """Interactive decoding."""
 
   def input_fn():
@@ -353,7 +353,7 @@ def input_fn():
     example = _interactive_input_tensor_to_features_dict(example, hparams)
     return example
 
-  result_iter = estimator.predict(input_fn)
+  result_iter = estimator.predict(input_fn, checkpoint_path=checkpoint_path)
   for result in result_iter:
     problem_idx = result["problem_choice"]
     is_image = False  # TODO(lukaszkaiser): find out from problem id / class.