Repository: OpenNMT/OpenNMT-tf
Pull request #820 (2021-03-25 11:50:40)

Fix loss normalization when using sentence weighting


Patch:
diff --git a/opennmt/runner.py b/opennmt/runner.py
index 06cd666f..fca7ca6f 100644
--- a/opennmt/runner.py
+++ b/opennmt/runner.py
@@ -144,6 +144,7 @@ class Runner(object):
                     max_batch_size=max_batch_size,
                     min_range=min_range,
                     num_devices=num_devices,
+                    scaling_factor=train_config.get("batch_size_autotune_scale", 0.8),
                     mixed_precision=self._mixed_precision,
                 )
 
@@ -501,7 +502,7 @@ def _auto_tune_batch_size(
     min_range,
     sample_iterations=5,
     num_devices=1,
-    scaling_factor=0.9,
+    scaling_factor=0.8,
     mixed_precision=False,
 ):
     """Find the largest token-based batch size that can be used with this
diff --git a/opennmt/utils/losses.py b/opennmt/utils/losses.py
index 64e96fbd..af5e89fc 100644
--- a/opennmt/utils/losses.py
+++ b/opennmt/utils/losses.py
@@ -51,23 +51,26 @@ def cross_entropy_sequence_loss(
     cross_entropy = _softmax_cross_entropy(logits, labels, label_smoothing, training)
     dtype = cross_entropy.dtype
 
-    if sequence_weight is not None:
-        cross_entropy *= tf.expand_dims(tf.cast(sequence_weight, dtype), 1)
+    shape = tf.shape(logits)
+    batch_size = shape[0]
+    max_time = shape[1]
 
-    if sequence_length is not None:
-        max_time = tf.shape(logits)[1]
-        mask = tf.sequence_mask(sequence_length, maxlen=max_time, dtype=dtype)
-        cross_entropy *= mask
-        loss_token_normalizer = tf.reduce_sum(mask)
-    else:
-        loss_token_normalizer = tf.cast(tf.size(labels), dtype)
+    if sequence_length is None:
+        sequence_length = tf.fill([batch_size], max_time)
 
-    loss = tf.reduce_sum(cross_entropy)
+    weight = tf.sequence_mask(sequence_length, maxlen=max_time, dtype=dtype)
+    if sequence_weight is not None:
+        sequence_weight = tf.cast(sequence_weight, dtype)
+        weight *= tf.expand_dims(sequence_weight, 1)
 
+    loss = tf.reduce_sum(cross_entropy * weight)
+
+    loss_token_normalizer = tf.reduce_sum(weight)
     if average_in_time or not training:
         loss_normalizer = loss_token_normalizer
+    elif sequence_weight is not None:
+        loss_normalizer = tf.reduce_sum(sequence_weight)
     else:
-        batch_size = tf.shape(logits)[0]
         loss_normalizer = tf.cast(batch_size, dtype)
 
     return loss, loss_normalizer, loss_token_normalizer
@@ -87,10 +90,16 @@ def cross_entropy_loss(logits, labels, label_smoothing=0.0, training=None, weigh
       The cumulated loss and the loss normalizer.
     """
     cross_entropy = _softmax_cross_entropy(logits, labels, label_smoothing, training)
+
     if weight is not None:
-        cross_entropy *= tf.cast(weight, cross_entropy.dtype)
+        weight = tf.cast(weight, cross_entropy.dtype)
+        cross_entropy *= weight
+        loss_normalizer = tf.reduce_sum(weight)
+    else:
+        batch_size = tf.shape(cross_entropy)[0]
+        loss_normalizer = tf.cast(batch_size, cross_entropy.dtype)
+
     loss = tf.reduce_sum(cross_entropy)
-    loss_normalizer = tf.cast(tf.shape(cross_entropy)[0], loss.dtype)
     return loss, loss_normalizer
 
 


Test patch:
diff --git a/opennmt/tests/losses_test.py b/opennmt/tests/losses_test.py
index 929634db..da1e1b6e 100644
--- a/opennmt/tests/losses_test.py
+++ b/opennmt/tests/losses_test.py
@@ -38,10 +38,11 @@ class LossesTest(tf.test.TestCase):
         labels = tf.constant([[2, 1, 0], [1, 0, 2]], dtype=tf.int32)
         lengths = tf.constant([2, 1], dtype=tf.int32)
 
-        loss, _, stats_norm = losses.cross_entropy_sequence_loss(
+        loss, train_norm, stats_norm = losses.cross_entropy_sequence_loss(
             logits, labels, sequence_length=lengths, training=True
         )
         self.assertNear(loss, 1.22118, 1e-5)
+        self.assertEqual(train_norm, 2)
         self.assertEqual(stats_norm, 3)
 
     def testWeightedAndMaskedCrossEntropySequenceLoss(self):
@@ -55,7 +56,7 @@ class LossesTest(tf.test.TestCase):
         lengths = tf.constant([3, 2], dtype=tf.int32)
         weights = tf.constant([0.6, 1.2])
 
-        loss, _, _ = losses.cross_entropy_sequence_loss(
+        loss, train_norm, stats_norm = losses.cross_entropy_sequence_loss(
             logits,
             labels,
             sequence_length=lengths,
@@ -63,6 +64,12 @@ class LossesTest(tf.test.TestCase):
             training=True,
         )
         self.assertNear(loss, 1.77306, 1e-5)
+        self.assertNear(train_norm, tf.reduce_sum(weights), 1e-5)
+        self.assertNear(
+            stats_norm,
+            tf.reduce_sum(tf.cast(lengths, tf.float32) * weights),
+            1e-5,
+        )
 
     @parameterized.expand(
         [