Fix transformer loss (#4270)

official/transformer/transformer_main.py

@@ -81,9 +81,12 @@ def model_fn(features, labels, mode, params):
     logits = output
 
     # Calculate model loss.
+    # xentropy contains the cross entropy loss of every nonpadding token in the
+    # targets.
     xentropy, weights = metrics.padded_cross_entropy_loss(
         logits, targets, params.label_smoothing, params.vocab_size)
-    loss = tf.reduce_sum(xentropy * weights) / tf.reduce_sum(weights)
+    # Compute the weighted mean of the cross entropy losses
+    loss = tf.reduce_sum(xentropy) / tf.reduce_sum(weights)
 
     # Save loss as named tensor that will be logged with the logging hook.
     tf.identity(loss, "cross_entropy")

official/transformer/utils/metrics.py

@@ -58,8 +58,8 @@ def padded_cross_entropy_loss(logits, labels, smoothing, vocab_size):
     smoothing: Label smoothing constant, used to determine the on and off values
     vocab_size: int size of the vocabulary
   Returns:
-    Returns a float32 tensor with shape
-      [batch_size, max(length_logits, length_labels)]
+    Returns the cross entropy loss and weight tensors: float32 tensors with
+      shape [batch_size, max(length_logits, length_labels)]
   """
   with tf.name_scope("loss", [logits, labels]):
     logits, labels = _pad_tensors_to_same_length(logits, labels)