Repository: OpenNMT/OpenNMT-tf (python)

Issue:
ValueError in guided alignment loss during training
Keras losses throw an exception when they are used within a distribution strategy scope and the reduction mode is unset.

```
ValueError: Please use `tf.keras.losses.Reduction.SUM` or `tf.keras.losses.Reduction.NONE` for loss reduction when losses are used with `tf.distribute.Strategy` outside of the built-in training loops. You can implement `tf.keras.losses.Reduction.SUM_OVER_BATCH_SIZE` using global batch size like:
```

Pull request:
Set reduction to support running GA loss under distribution strategy
Fixes #526.

Patch:
diff --git a/opennmt/utils/losses.py b/opennmt/utils/losses.py
index 267d0be6..67e8acd0 100644
--- a/opennmt/utils/losses.py
+++ b/opennmt/utils/losses.py
@@ -101,9 +101,9 @@ def guided_alignment_cost(attention_probs,
     ValueError: if :obj:`cost_type` is invalid.
   """
   if cost_type == "ce":
-    loss = tf.keras.losses.CategoricalCrossentropy()
+    loss = tf.keras.losses.CategoricalCrossentropy(reduction=tf.keras.losses.Reduction.SUM)
   elif cost_type == "mse":
-    loss = tf.keras.losses.MeanSquaredError()
+    loss = tf.keras.losses.MeanSquaredError(reduction=tf.keras.losses.Reduction.SUM)
   else:
     raise ValueError("invalid guided alignment cost: %s" % cost_type)
 
@@ -113,13 +113,16 @@ def guided_alignment_cost(attention_probs,
         maxlen=tf.shape(attention_probs)[1],
         dtype=attention_probs.dtype)
     sample_weight = tf.expand_dims(sample_weight, -1)
+    normalizer = tf.reduce_sum(sequence_length)
   else:
     sample_weight = None
+    normalizer = tf.size(attention_probs)
 
   cost = loss(
       gold_alignment,
       attention_probs,
       sample_weight=sample_weight)
+  cost /= tf.cast(normalizer, cost.dtype)
   return weight * cost
 
 def regularization_penalty(regularization_type, scale, weights):


Test patch:
diff --git a/opennmt/tests/losses_test.py b/opennmt/tests/losses_test.py
index 61d74771..6a95ee97 100644
--- a/opennmt/tests/losses_test.py
+++ b/opennmt/tests/losses_test.py
@@ -30,6 +30,26 @@ class LossesTest(tf.test.TestCase):
     with self.assertRaises(ValueError):
       losses.regularization_penalty("l1_l2", 1e-4, [])
 
+  @parameterized.expand([
+      ["ce", False],
+      ["mse", False],
+      ["mse", True],
+  ])
+  def testGuidedAlignmentCostUnderDistributionStrategy(self, cost_type, with_length):
+    strategy = tf.distribute.MirroredStrategy(devices=["/cpu:0"])
+    attention_probs = tf.random.uniform([2, 5, 6])
+    gold_alignment = tf.random.uniform([2, 5, 6])
+    if with_length:
+      sequence_length = tf.constant([4, 5], dtype=tf.int32)
+    else:
+      sequence_length = None
+    with strategy.scope():
+      losses.guided_alignment_cost(
+          attention_probs,
+          gold_alignment,
+          sequence_length=sequence_length,
+          cost_type=cost_type)
+
 
 if __name__ == "__main__":
   tf.test.main()