Repository: OpenNMT/OpenNMT-tf

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:
```


Discussion:
I've happen to reproduce this issue.

OpenNMT-tf v2.1.0
Python v3.5.2
tensorflow-gpu v2.0.0

I'm trying to train a Transformer model with guided alignment. I've used SentencePiece to tokenize EN-RU corpora with BPE and then I've ran `fast_align` with `grow-diag-final-and`.

```
▁I ▁found ▁them ▁in ▁the ▁garbage ▁.
▁But ▁look ▁how ▁many ▁there ▁are ▁.
▁Well ▁, ▁let ▁me ▁take ▁a ▁look ▁at ▁' ▁em ▁.
```

```
▁На шёл ▁на ▁помо й ке ▁.
▁Зато ▁смотри ▁как ▁их ▁много ▁.
▁Ладно ▁, ▁дай ▁взгля ну ▁поближе ▁.
```

```
1-0 3-2 5-3 6-6
0-0 1-1 2-2 3-3 6-5
0-0 1-1 2-2 6-3 10-6
```

Here is actual basic configuration that I use:

```
model_dir: toy_enru_transformer_withalign

data:
  train_features_file: data/train.en
  train_labels_file: data/train.ru
  train_alignments: data/train.align
  eval_features_file: data/valid.en
  eval_labels_file: data/valid.ru
  source_vocabulary: data/mdl-en.vocab
  target_vocabulary: data/mdl-ru.vocab

params:
  guided_alignment_type: "ce"
  guided_alignment_weight: 1.0

train:
  batch_type: tokens
  save_checkpoints_steps: 5000
  keep_checkpoint_max: 8

eval:
  external_evaluators: BLEU

infer:
  batch_size: 32
```

Running this command produces following output:

$ `onmt-main --model_type Transformer --config config/toy.yml --auto_config train --with_eval`

```
one NUMA node, so returning NUMA node zero
2019-10-18 09:59:04.801496: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1304] Created TensorFlow device (/device:GPU:7 with 15021 MB memory) -> physical GPU (device: 7, name: Tesla V100-SXM2-16GB, pci bus id: 0000:00:1e.0, compute capability: 7.0)
2019-10-18 09:59:04.962543: I tensorflow/stream_executor/platform/default/dso_loader.cc:44] Successfully opened dynamic library libcublas.so.10.0
WARNING:tensorflow:From /usr/local/lib/python3.5/dist-packages/tensorflow_core/python/framework/constant_op.py:253: _EagerTensorBase.cpu (from tensorflow.python.framework.ops) is deprecated and will be removed in a future version.
Instructions for updating:
Use tf.identity instead.
INFO:tensorflow:Saved checkpoint toy_enru_transformer_withalign/ckpt-0
INFO:tensorflow:Error reported to Coordinator: in converted code:
    relative to /usr/local/lib/python3.5/dist-packages:

    opennmt/training.py:87 _accumulate_gradients  *
        loss = self._model.compute_loss(outputs, target, training=True)
    opennmt/models/sequence_to_sequence.py:333 compute_loss  *
        loss += losses.guided_alignment_cost(
    opennmt/utils/losses.py:119 guided_alignment_cost  *
        cost = loss(
    tensorflow_core/python/keras/losses.py:128 __call__
        losses, sample_weight, reduction=self._get_reduction())
    tensorflow_core/python/keras/losses.py:162 _get_reduction
        'Please use `tf.keras.losses.Reduction.SUM` or '

    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:
    ```
    with strategy.scope():
        loss_obj = tf.keras.losses.CategoricalCrossentropy(reduction=tf.keras.losses.reduction.NONE)
    ....
        loss = tf.reduce_sum(loss_obj(labels, predictions)) * (1. / global_batch_size)
    ```
    Please see https://www.tensorflow.org/alpha/tutorials/distribute/training_loops for more details.
Traceback (most recent call last):
  File "/usr/local/lib/python3.5/dist-packages/tensorflow_core/python/training/coordinator.py", line 297, in stop_on_exception
    yield
  File "/usr/local/lib/python3.5/dist-packages/tensorflow_core/python/distribute/mirrored_strategy.py", line 879, in run
    self.main_result = self.main_fn(*self.main_args, **self.main_kwargs)
  File "/usr/local/lib/python3.5/dist-packages/tensorflow_core/python/autograph/impl/api.py", line 237, in wrapper
    raise e.ag_error_metadata.to_exception(e)
ValueError: in converted code:
    relative to /usr/local/lib/python3.5/dist-packages:

    opennmt/training.py:87 _accumulate_gradients  *
        loss = self._model.compute_loss(outputs, target, training=True)
    opennmt/models/sequence_to_sequence.py:333 compute_loss  *
        loss += losses.guided_alignment_cost(
    opennmt/utils/losses.py:119 guided_alignment_cost  *
        cost = loss(
    tensorflow_core/python/keras/losses.py:128 __call__
        losses, sample_weight, reduction=self._get_reduction())
    tensorflow_core/python/keras/losses.py:162 _get_reduction
        'Please use `tf.keras.losses.Reduction.SUM` or '

    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:
    ```
    with strategy.scope():
        loss_obj = tf.keras.losses.CategoricalCrossentropy(reduction=tf.keras.losses.reduction.NONE)
    ....
        loss = tf.reduce_sum(loss_obj(labels, predictions)) * (1. / global_batch_size)
    ```
    Please see https://www.tensorflow.org/alpha/tutorials/distribute/training_loops for more details.

Traceback (most recent call last):
  File "/usr/local/bin/onmt-main", line 8, in <module>
    sys.exit(main())
  File "/usr/local/lib/python3.5/dist-packages/opennmt/bin/main.py", line 189, in main
    checkpoint_path=args.checkpoint_path)
  File "/usr/local/lib/python3.5/dist-packages/opennmt/runner.py", line 205, in train
    export_on_best=eval_config.get("export_on_best"))
  File "/usr/local/lib/python3.5/dist-packages/opennmt/training.py", line 146, in __call__
    for i, (loss, num_words) in enumerate(_forward()):  # pylint: disable=no-value-for-parameter
  File "/usr/local/lib/python3.5/dist-packages/opennmt/data/dataset.py", line 433, in _fun
    outputs = _tf_fun()
  File "/usr/local/lib/python3.5/dist-packages/tensorflow_core/python/eager/def_function.py", line 457, in __call__
    result = self._call(*args, **kwds)
  File "/usr/local/lib/python3.5/dist-packages/tensorflow_core/python/eager/def_function.py", line 503, in _call
    self._initialize(args, kwds, add_initializers_to=initializer_map)
  File "/usr/local/lib/python3.5/dist-packages/tensorflow_core/python/eager/def_function.py", line 408, in _initialize
    *args, **kwds))
  File "/usr/local/lib/python3.5/dist-packages/tensorflow_core/python/eager/function.py", line 1848, in _get_concrete_function_internal_garbage_collected
    graph_function, _, _ = self._maybe_define_function(args, kwargs)
  File "/usr/local/lib/python3.5/dist-packages/tensorflow_core/python/eager/function.py", line 2150, in _maybe_define_function
    graph_function = self._create_graph_function(args, kwargs)
  File "/usr/local/lib/python3.5/dist-packages/tensorflow_core/python/eager/function.py", line 2041, in _create_graph_function
    capture_by_value=self._capture_by_value),
  File "/usr/local/lib/python3.5/dist-packages/tensorflow_core/python/framework/func_graph.py", line 915, in func_graph_from_py_func
    func_outputs = python_func(*func_args, **func_kwargs)
  File "/usr/local/lib/python3.5/dist-packages/tensorflow_core/python/eager/def_function.py", line 358, in wrapped_fn
    return weak_wrapped_fn().__wrapped__(*args, **kwds)
  File "/usr/local/lib/python3.5/dist-packages/tensorflow_core/python/framework/func_graph.py", line 905, in wrapper
    raise e.ag_error_metadata.to_exception(e)
ValueError: in converted code:
    relative to /usr/local/lib/python3.5/dist-packages:

    opennmt/data/dataset.py:429 _tf_fun  *
        return func(lambda: next(iterator))
    opennmt/training.py:122 _forward  *
        per_replica_loss, per_replica_words = self._strategy.experimental_run_v2(
    tensorflow_core/python/distribute/distribute_lib.py:760 experimental_run_v2
        return self._extended.call_for_each_replica(fn, args=args, kwargs=kwargs)
    opennmt/training.py:87 _accumulate_gradients  *
        loss = self._model.compute_loss(outputs, target, training=True)
    opennmt/models/sequence_to_sequence.py:333 compute_loss  *
        loss += losses.guided_alignment_cost(
    opennmt/utils/losses.py:119 guided_alignment_cost  *
        cost = loss(
    tensorflow_core/python/keras/losses.py:128 __call__
        losses, sample_weight, reduction=self._get_reduction())
    tensorflow_core/python/keras/losses.py:162 _get_reduction
        'Please use `tf.keras.losses.Reduction.SUM` or '

    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:
    ```
    with strategy.scope():
        loss_obj = tf.keras.losses.CategoricalCrossentropy(reduction=tf.keras.losses.reduction.NONE)
    ....
        loss = tf.reduce_sum(loss_obj(labels, predictions)) * (1. / global_batch_size)
    ```
    Please see https://www.tensorflow.org/alpha/tutorials/distribute/training_loops for more details.
```

Please could you elaborate on that. Maybe there is a workaround to train Transformer with guided alignment on > v2.0.0?
This just requires a fix in the code. Transformer with guided alignment works by itself but not when used with the training utilities that run the model within a [distribution strategy](https://www.tensorflow.org/api_docs/python/tf/distribute).
Ok. I'm really waiting forward for this to be fixed.
As my understanding goes I can pass a reduction of `SUM` to loss in `losses.py` and it should work correctly. Can this actually work?

Resolving patch:
diff --git a/opennmt/utils/losses.py b/opennmt/utils/losses.py
--- 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
--- a/opennmt/tests/losses_test.py
+++ b/opennmt/tests/losses_test.py
@@ -30,6 +30,26 @@ def testRegulaizationMissingScaleValue(self):
     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()