TransformerScorer model to only score targets on infer

PiperOrigin-RevId: 192812089

Author: Ryan Sepassi

tensor2tensor/models/transformer.py

@@ -535,6 +535,59 @@ def is_not_finished(i, finished, *_):
   return {"outputs": decoded_ids, "scores": scores}
 
 
+@registry.register_model
+class TransformerScorer(Transformer):
+  """Transformer model, but only scores in PREDICT mode.
+
+  Checkpoints between Transformer and TransformerScorer are interchangeable.
+  """
+
+  def __init__(self, *args, **kwargs):
+    super(TransformerScorer, self).__init__(*args, **kwargs)
+    self._name = "transformer"
+    self._base_name = "transformer"
+
+  def infer(self,
+            features=None,
+            decode_length=50,
+            beam_size=1,
+            top_beams=1,
+            alpha=0.0):
+    """Returns the targets and their log probabilities."""
+    del decode_length, beam_size, top_beams, alpha
+    assert features is not None
+
+    # Run the model
+    self.hparams.force_full_predict = True
+    with tf.variable_scope(self.name):
+      logits, _ = self.model_fn(features)
+    assert len(logits.shape) == 5  # [batch, time, 1, 1, vocab]
+    logits = tf.squeeze(logits, [2, 3])
+
+    # Compute the log probabilities
+    log_probs = beam_search.log_prob_from_logits(logits)
+
+    # Slice out the log_probs of the targets
+    targets = features["targets"]
+    assert len(targets.shape) == 4  # [batch, time, 1, 1]
+    targets = tf.squeeze(targets, [2, 3])
+    batch_size, timesteps = common_layers.shape_list(targets)
+    vocab_size = common_layers.shape_list(log_probs)[-1]
+    flat_targets = tf.reshape(targets, [batch_size * timesteps])
+    flat_log_probs = tf.reshape(log_probs, [batch_size * timesteps, vocab_size])
+    flat_indices = tf.stack(
+        [tf.range(tf.to_int64(batch_size) * tf.to_int64(timesteps)),
+         tf.to_int64(flat_targets)], axis=1)
+    log_probs = tf.reshape(
+        tf.gather_nd(flat_log_probs, flat_indices),
+        [batch_size, timesteps])
+
+    # Sum over time to get the log_prob of the sequence
+    scores = tf.reduce_sum(log_probs, axis=1)
+
+    return {"outputs": targets, "scores": scores}
+
+
 @registry.register_model
 class TransformerEncoder(t2t_model.T2TModel):
   """Transformer, encoder only."""

tensor2tensor/models/transformer_test.py

@@ -37,7 +37,10 @@
 
 class TransformerTest(tf.test.TestCase):
 
-  def getModel(self, hparams, mode=tf.estimator.ModeKeys.TRAIN, has_input=True):
+  def getModel(self, hparams=None, mode=tf.estimator.ModeKeys.TRAIN,
+               has_input=True, model_cls=transformer.Transformer):
+    if hparams is None:
+      hparams = transformer.transformer_tiny()
     hparams.hidden_size = 8
     hparams.filter_size = 32
     hparams.num_heads = 1
@@ -58,7 +61,7 @@ def getModel(self, hparams, mode=tf.estimator.ModeKeys.TRAIN, has_input=True):
         "target_space_id": tf.constant(1, dtype=tf.int32)
     }
 
-    return transformer.Transformer(hparams, mode, p_hparams), features
+    return model_cls(hparams, mode, p_hparams), features
 
   def testTransformer(self):
     model, features = self.getModel(transformer.transformer_small())
@@ -240,5 +243,47 @@ def testTransformerWithEncoderDecoderAttentionLoss(self):
     self.assertEqual(res.shape, ())
 
 
+class TransformerScorerTest(TransformerTest):
+
+  def testReturnsScores(self):
+    model, features = self.getModel(
+        mode=tf.estimator.ModeKeys.PREDICT,
+        model_cls=transformer.TransformerScorer)
+    infer_out = model.infer(features)
+    self.assertTrue("outputs" in infer_out)
+    self.assertTrue("scores" in infer_out)
+
+    with self.test_session() as session:
+      session.run(tf.global_variables_initializer())
+      infer_out = session.run(infer_out)
+      self.assertEqual((BATCH_SIZE,), infer_out["scores"].shape)
+      self.assertEqual((BATCH_SIZE, TARGET_LENGTH), infer_out["outputs"].shape)
+
+  def testVarNames(self):
+    with tf.Graph().as_default():
+      model, features = self.getModel(
+          mode=tf.estimator.ModeKeys.PREDICT,
+          model_cls=transformer.TransformerScorer)
+      _ = model.infer(features)
+      scorer_vars = [v.name for v in tf.global_variables()]
+
+    with tf.Graph().as_default():
+      model, features = self.getModel(
+          mode=tf.estimator.ModeKeys.EVAL,
+          model_cls=transformer.TransformerScorer)
+      _ = model(features)
+      scorer_eval_vars = [v.name for v in tf.global_variables()]
+
+    with tf.Graph().as_default():
+      model, features = self.getModel(
+          mode=tf.estimator.ModeKeys.EVAL,
+          model_cls=transformer.Transformer)
+      _ = model(features)
+      transformer_vars = [v.name for v in tf.global_variables()]
+
+    self.assertEqual(sorted(scorer_vars), sorted(transformer_vars))
+    self.assertEqual(sorted(scorer_eval_vars), sorted(transformer_vars))
+
+
 if __name__ == "__main__":
   tf.test.main()