Repository: huggingface/transformers

Issue:
Potential incorrect loss calculation for TFTokenClassification in TFTrainer
## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->
     
- `transformers` version: 3.1.0
- Platform: Linux-4.15.0-115-generic-x86_64-with-debian-buster-sid
- Python version: 3.6.7
- PyTorch version (GPU?): 1.5.1+cpu (False)
- Tensorflow version (GPU?): 2.2.0 (False)
- Using GPU in script?: No
- Using distributed or parallel set-up in script?: No


### Who can help
<!-- Your issue will be replied to more quickly if you can figure out the right person to tag with @
 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 -->
 
 Trainer: @sgugger
 tensorflow: @jplu
 examples/token-classification: @stefan-it

Mostly for @jplu, potentially for @stefan-it (because the workaround I have in mind requires a bit change in the token classification dataset).

## Information

The problem arises when using:
* [x] The official example scripts:
    The involved scripts are:
    - https://github.com/huggingface/transformers/blob/master/src/transformers/trainer_tf.py
    - https://github.com/huggingface/transformers/blob/master/examples/token-classification/run_tf_ner.py

  However, in order to demonstrate the issue in a more clear way, I use a minimal example which doesn't use directly these two scripts. See the description and code snippet below.

The tasks I am working on is:
* [x] Official token classification task in TensorFlow
  
## Description

In [trainer_tf.py](https://github.com/huggingface/transformers/blob/master/src/transformers/trainer_tf.py#L595), the loss calculation is calculated from `per_example_loss` divided by `total_train_batch_size`.

        per_example_loss, _ = self.run_model(features, labels, True)
        scaled_loss = per_example_loss / self.total_train_batch_size

Here `total_train_batch_size` is the size of a whole batch that will be distributed to (potentially) different replicas and optionally consisting of several smaller batches for accumulation steps.

For sentence level tasks, where each example (i.e. sentence) corresponds to a label (for example, sentence classification), the above loss calculation is correct.

However, for token level tasks like token classification, the the above loss seems incorrect to me. For such tasks, the loss should be the per example losses **divided by the number of real tokens involved in the batch**.

In [utils_ner](https://github.com/huggingface/transformers/blob/master/examples/token-classification/utils_ner.py#L75), `convert_examples_to_features` set labels to `-100` for padding tokens and other special tokens (`[CLS]`, `[SEP]`, etc), which are the places to be ignored for loss calculation. Therefore, the loss calculation should be the per example losses **divided by the number of labels that are not -100 in the \*_batch_\***.

By **\*_batch_\***, it should be careful that it is not the batch received by a single replica, and neither the smaller batch in a single accumulation step. It means `the whole batch that will be distributed to (potentially) different replicas and optionally consisting of several smaller batches for accumulation steps.` More precisely, it means a batch passed to [distributed_training_steps()](https://github.com/huggingface/transformers/blob/master/src/transformers/trainer_tf.py#L651) - for the same reason as we divide per example losses by `total_train_batch_size` for sentence level tasks, rather than dividing it by the size of batch received by a single replica.

In order to calculate the correct loss values, we have to pass the global information - the number of labels that are not `-100` in a `global batch` to each replica. I don't know a clean way to do it, but for my own personal projects, I inject this extra information into global batches as a constant, and each replica receiving a distributed smaller batch will have this information to calculate the correct scaled losses.

(I have a notebook showing how to perform it, if you want to look it, let me know.)

## Code Snippets
<!-- If you have code snippets, error messages, stack traces please provide them here as well.
     Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting
     Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.-->

Here is a minimal example to demonstrate the issue. 

Here, we have only one real example (sentence) and `n_empty_string` empty sentences.
Each empty sentence will give only [CLS], [SEP] and [PAD] tokens that will be ignored for token classification.

    import os
    os.environ['TF_DETERMINISTIC_OPS'] = '1'
    
    SEED = 42
    name = 'distilbert-base-uncased'
    seq_len = 8
    num_labels = 2
    n_empty_string = 10
    
    import tensorflow as tf
    tf.random.set_seed(SEED)
    
    strategy = tf.distribute.OneDeviceStrategy(device="/cpu:0")
    
    from transformers import TFTrainer, AutoConfig, AutoTokenizer, TFAutoModelForTokenClassification
    from transformers.training_args_tf import TFTrainingArguments
    
    text = [
        'My dog is cute'
    ]
    text.extend([''] * n_empty_string)
    n_examples = len(text)
    
    config = AutoConfig.from_pretrained(
        name,
        num_labels=num_labels
    )
    
    tokenizer = AutoTokenizer.from_pretrained(name)
    
    model = TFAutoModelForTokenClassification.from_pretrained(
        name
    )
    training_args = TFTrainingArguments(
        output_dir='./tmp/',
        per_device_train_batch_size=n_examples,
        gradient_accumulation_steps=1,
        seed=SEED
    )
    
    # Initialize our Trainer
    trainer = TFTrainer(
        model=model,
        args=training_args,
        train_dataset=None,
        eval_dataset=None,
        compute_metrics=None
    )
    trainer.total_train_batch_size = strategy.num_replicas_in_sync \
                                     * training_args.per_device_train_batch_size \
                                     * training_args.gradient_accumulation_steps
    trainer.train_loss = tf.keras.metrics.Sum()
    
    features = tokenizer.batch_encode_plus(text, max_length=seq_len, padding='max_length', return_tensors='tf')
    # Set all labels to `1`, except for special tokens: cls/sep/pad, where the labels are `-100`.
    labels = tf.constant(1, shape=[n_examples, seq_len])
    for token_id in [tokenizer.pad_token_id] + tokenizer.all_special_ids:
        labels = labels * tf.cast(features['input_ids'] != token_id, dtype=tf.int32) + \
                 -100 * tf.cast(features['input_ids'] == token_id, dtype=tf.int32)
    
    # Only the first example `features[0]` has real tokens, the other examples have only [PAD].
    print(features['input_ids'])
    
    # Only the first example has labels that won't be ignored.
    print(labels)
    
    # Copy from:
    #     https://github.com/huggingface/transformers/blob/master/src/transformers/trainer_tf.py#L601
    per_example_loss, _ = trainer.run_model(features, labels, True)
    scaled_loss = per_example_loss / trainer.total_train_batch_size
    
    print(scaled_loss)

## Expected behavior

<!-- A clear and concise description of what you would expect to happen. -->
When `n_empty_string = 0`, we get `scaled_loss`

    tf.Tensor([0.56047076 0.46507886 0.51456743 0.50131255], shape=(4,), dtype=float32)

When `n_empty_string = 9`, we get `scaled_loss`

    tf.Tensor([0.05604707 0.04650789 0.05145674 0.05013125], shape=(4,), dtype=float32)

However, in both case, we should get the same value, which should be

    tf.Tensor([0.56047076 0.46507886 0.51456743 0.50131255], shape=(4,), dtype=float32)


Discussion:
Hello @chiapas!

Thanks for having investigating this! Have you checked the way we compute the loss for Token classification? Right here https://github.com/huggingface/transformers/blob/master/src/transformers/modeling_tf_utils.py#L161

We are already ignoring all the tokens that have -100 as label. So if there is an issue it might come from somewhere else.

Hi, @jplu,

Yes, I checked that. But the issue in this bug report is not about ignoring -100 or not. The problem is that, the loss is calculated from the per example losses, then divided by `total_train_batch_size`. However, for token level tasks, it should be divided by `the number of actual tokens (i.e. tokens not ignored) in the global batch (i.e. the batch that having the size total_train_batch_size)`.
I don't get what you mean by the number of actual token?

You mean the number of batches that contain actual tokens, no? In your example, just 1? 
If you want a support for the argument about the denominator I claimed, we can look at the pytorch implementation in [DistilBert](https://github.com/huggingface/transformers/blob/master/src/transformers/modeling_distilbert.py#L820) about the loss for token classification, you will see

    loss_fct = CrossEntropyLoss() 

and 

    loss = loss_fct(active_logits, active_labels)

And from [torch's doc](https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html#crossentropyloss), the default reduction is `mean`. So this corresponds to the per example losses (-100 ignored) divided by the number of actual tokens (i.e. -100 ignored again).




> 
> 
> I don't get what you mean by the number of actual token?
> 
> You mean the number of batches that contain actual tokens, no? In your example, just 1?

For any global batch (that has size `total_train_batch_size`), after distributed to replicas, we compute the per example losses on the smaller batches (where the tokens with label -100 are ignored). Then in the current implementation, this per example losses are divided by `total_train_batch_size`.

By the actual tokens, it means `the tokens in a global batch with labels != -100`. And what I said is that, the per example losses should be divided by the number of tokens with label != -100 in that global batch.

This number of `actual` tokens will be varied for different global batch however.

In the example in the code snippet, they are only `4` actual tokens. By when `n_empty_string = 9`, the current implementation divided the per example losses by ` 10 (1 + 9 dummy sentences)`. 

So basically what you propose is to return a mean reduction in the `call` function instead of the effective per example loss? Something like to replace:

```
loss = None if labels is None else self.compute_loss(labels, logits)
```
By
```
loss = None if labels is None else tf.reduce_mean(self.compute_loss(labels, logits))
```

And then dividing the result per the number of replica:
```
per_example_loss, _ = self.run_model(features, labels, True)
scaled_loss = per_example_loss / self.args.n_replicas
```

Otherwise, please show me you notebook because I still don't get it.
@jplu ,

I will explain a bit more and also show my notebook later. But no, I am not suggesting using

    tf.reduce_mean(self.compute_loss(labels, logits))  # If we do so, the average occurs on the small batch in each replica.

I mentioned the pytorch version just to show that the loss should be averaged over the tokens, not over the sentences in the batch. However, the average shouldn't be over the small batches received by each replica, it should be over the global batches.
@jplu 

If you want  to look the code directly, here is my (kaggle) notebook

[Masked, My Dear Watson - MLM with TPU](https://www.kaggle.com/yihdarshieh/masked-my-dear-watson-mlm-with-tpu#MLM-loss-calculation). Please check `def mlm_fine_tune_step(batch):` just below that markdown cell,  which has

    loss_mlm = loss_fn(
        labels_at_masked_tokens,
        logits_at_masked_tokens
    )

    # divide the number of masked tokens in the global batch, i.e. the whole batch that is distributed to different replicas.
    loss_mlm = loss_mlm / tf.cast(nb_tokens_masked[0], dtype=tf.float32)

If you prefer, I can work on this and make a PR. But I think it is better for us to agree that the loss calculation should be corrected. So I will try to explain:

1.  For token level tasks, the loss value is the per example (and here, example = tokens) losses in a batch, divided by the number of tokens in that batch.
2.  If we have tokens being ignored for loss calculation, the denominator above becomes the number of tokens not ignored in that batch.
3. By `batch`, it is the whole set used for 1 parameter update by gradients - which is called a `global batch`.
4. Since we use distributed strategy, and optionally gradient accumulation, while `training_step()` processes a batch, it is a small batch (i.e. a batch for `only 1 gradient accumulation step` on a `single replica`). However, the denominator in step `1.` or `2.` should be the `the number of token, not being ignored, in a global batch`, despite the per example losses are still based on the small batch received by a replica.
5. Since gradient accumulation will add the gradients, and distributed strategy will sync across replicas by summing the gradients before applying gradients, the above steps will give us the `averaged losses over the tokens (not ignored) in a global batch`.

Hope it make things a bit clear.
OK now with an example and the explanation I got it. Thank you very much!

I prefer you do a PR and then you get the credit of this fix :) And if you can tag me as reviewer I will be able to help you if needed, as there is certainly a nicer way to do. Maybe with a class field?

Thanks again, waiting your PR ^^

Resolving patch:
diff --git a/src/transformers/trainer_tf.py b/src/transformers/trainer_tf.py
--- a/src/transformers/trainer_tf.py
+++ b/src/transformers/trainer_tf.py
@@ -9,6 +9,7 @@
 import numpy as np
 import tensorflow as tf
 from packaging.version import parse
+from tensorflow.python.distribute.values import PerReplica
 
 from .integrations import is_comet_available, is_wandb_available
 from .modeling_tf_utils import TFPreTrainedModel
@@ -363,7 +364,7 @@ def prediction_loop(
         else:
             metrics = {}
 
-        metrics["eval_loss"] = self.eval_loss.result().numpy() / (steps * self.args.eval_batch_size)
+        metrics["eval_loss"] = self.eval_loss.result().numpy() / steps
 
         for key in list(metrics.keys()):
             if not key.startswith("eval_"):
@@ -441,21 +442,28 @@ def evaluate(self, eval_dataset: Optional[tf.data.Dataset] = None) -> Dict[str,
 
         return output.metrics
 
-    def prediction_step(self, features: tf.Tensor, labels: tf.Tensor) -> tf.Tensor:
+    def prediction_step(
+        self, features: tf.Tensor, labels: tf.Tensor, nb_instances_in_global_batch: tf.Tensor
+    ) -> tf.Tensor:
         """
         Compute the prediction on features and update the loss with labels.
 
         Subclass and override to inject some custom behavior.
         """
         per_example_loss, logits = self.run_model(features, labels, False)
+        scaled_loss = per_example_loss / tf.cast(nb_instances_in_global_batch, dtype=per_example_loss.dtype)
 
-        self.eval_loss.update_state(per_example_loss)
+        self.eval_loss.update_state(scaled_loss)
 
         return logits
 
     @tf.function
     def distributed_prediction_steps(self, batch):
-        logits = self.args.strategy.run(self.prediction_step, batch)
+
+        nb_instances_in_batch = self._compute_nb_instances(batch)
+        inputs = self._get_step_inputs(batch, nb_instances_in_batch)
+
+        logits = self.args.strategy.run(self.prediction_step, inputs)
 
         return logits
 
@@ -542,7 +550,7 @@ def train(self) -> None:
 
                     self.distributed_training_steps(batch)
 
-                    training_loss = self.train_loss.result() / ((step + 1) * self.total_train_batch_size)
+                    training_loss = self.train_loss.result() / (step + 1)
 
                     if self.args.debug:
                         logs = {}
@@ -592,14 +600,14 @@ def train(self) -> None:
             # Clean the state at the end of training
             delattr(self, "_past")
 
-    def training_step(self, features, labels):
+    def training_step(self, features, labels, nb_instances_in_global_batch):
         """
         Perform a training step on features and labels.
 
         Subclass and override to inject some custom behavior.
         """
         per_example_loss, _ = self.run_model(features, labels, True)
-        scaled_loss = per_example_loss / self.total_train_batch_size
+        scaled_loss = per_example_loss / tf.cast(nb_instances_in_global_batch, dtype=per_example_loss.dtype)
         gradients = tf.gradients(scaled_loss, self.model.trainable_variables)
         gradients = [
             g if g is not None else tf.zeros_like(v) for g, v in zip(gradients, self.model.trainable_variables)
@@ -608,14 +616,14 @@ def training_step(self, features, labels):
         if self.args.gradient_accumulation_steps > 1:
             self.gradient_accumulator(gradients)
 
-        self.train_loss.update_state(per_example_loss)
+        self.train_loss.update_state(scaled_loss)
 
         if self.args.gradient_accumulation_steps == 1:
             return gradients
 
-    def apply_gradients(self, features, labels):
+    def apply_gradients(self, features, labels, nb_instances_in_global_batch):
         if self.args.gradient_accumulation_steps == 1:
-            gradients = self.training_step(features, labels)
+            gradients = self.training_step(features, labels, nb_instances_in_global_batch)
 
             self.optimizer.apply_gradients(list(zip(gradients, self.model.trainable_variables)))
         else:
@@ -625,7 +633,7 @@ def apply_gradients(self, features, labels):
                 }
                 reduced_labels = labels[: self.args.train_batch_size // self.args.n_replicas]
 
-                self.training_step(reduced_features, reduced_labels)
+                self.training_step(reduced_features, reduced_labels, nb_instances_in_global_batch)
 
                 features = {
                     k: tf.concat(
@@ -650,7 +658,35 @@ def apply_gradients(self, features, labels):
     @tf.function
     def distributed_training_steps(self, batch):
         with self.args.strategy.scope():
-            self.args.strategy.run(self.apply_gradients, batch)
+
+            nb_instances_in_batch = self._compute_nb_instances(batch)
+            inputs = self._get_step_inputs(batch, nb_instances_in_batch)
+
+            self.args.strategy.run(self.apply_gradients, inputs)
+
+    @staticmethod
+    def _compute_nb_instances(batch):
+
+        labels = batch[-1]
+        if isinstance(labels, PerReplica):
+            labels = tf.concat(labels.values, axis=0)
+
+        nb_instances = tf.reduce_sum(tf.cast(labels != -100, dtype=tf.int32))
+
+        return nb_instances
+
+    @staticmethod
+    def _get_step_inputs(batch, nb_instances):
+
+        features, labels = batch
+
+        if isinstance(labels, PerReplica):
+            # need to make a `PerReplica` objects for ``nb_instances``
+            nb_instances = PerReplica([nb_instances] * len(labels.values))
+
+        step_inputs = (features, labels, nb_instances)
+
+        return step_inputs
 
     def run_model(self, features, labels, training):
         """


Test patch: