Fix TFTrainer prediction output
Author: janinaj
Pull request: yes

# What does this PR do?

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

This PR fixes two issues:

1) The prediction output (specifically prediction_loop) of TFTrainer does not match the dataset cardinality. If the number of examples is divisible by eval_batch_size, the first batch is predicted twice. Else, the first n examples, where n = eval_batch_size - num_examples % eval_batch_size, are predicted twice.

This results in an output shape that is different from dataset cardinality. This also causes the output of evaluate(), including eval_loss, to be incorrect (e.g. loss is computed twice for the first few examples).

2) The evaluation loss only works the first time it is computed. Subsequent computations result to 0. Below is a sample output when an evaluation strategy is set during training:

[INFO|trainer_tf.py:398] 2021-01-18 01:34:58,856 >> {**'eval_loss': 0.6290212145038679**, 'eval_acc': 0.6875, ... 'step': 10}
[INFO|trainer_tf.py:398] 2021-01-18 01:35:10,852 >> {**'eval_loss': 0.0**, 'eval_acc': 0.6875, ..., 'step': 20}

<!-- Remove if not applicable -->


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md#start-contributing-pull-requests),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/master/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/master/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?

## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors which may be interested in your PR.

 tensorflow: @jplu
Trainer: @sgugger

<!-- Your PR 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**.
 Please tag fewer than 3 people.

 albert, bert, XLM: @LysandreJik
 GPT2: @LysandreJik, @patrickvonplaten
 tokenizers: @mfuntowicz
 Trainer: @sgugger
 Benchmarks: @patrickvonplaten
 Model Cards: @julien-c
 examples/distillation: @VictorSanh
 nlp datasets: [different repo](https://github.com/huggingface/nlp)
 rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)
 Text Generation: @patrickvonplaten, @TevenLeScao
 Blenderbot, Bart, Marian, Pegasus: @patrickvonplaten
 T5: @patrickvonplaten
 Rag: @patrickvonplaten, @lhoestq
 EncoderDecoder: @patrickvonplaten
 Longformer, Reformer: @patrickvonplaten
 TransfoXL, XLNet: @TevenLeScao, @patrickvonplaten
 examples/seq2seq: @patil-suraj
 examples/bert-loses-patience: @JetRunner
 tensorflow: @jplu
 examples/token-classification: @stefan-it
 documentation: @sgugger
 FSMT: @stas00
 -->


Comments:
Sorry for the confusion. Here is a detailed description of the first issue:

If the number of examples is divisible by eval_batch_size, the first batch is predicted twice. Say we have the following examples (n=8): A, B, C, D, E, F, G, H and eval_batch_size = 4. This will create 2 batches: (A,B,C,D) and (E,F,G,H). The current implementation of prediction_loop() will run batch (A,B,C,D), (E,F,G,H) and (A,B,C,D) again. This produces an output with a shape of 12 instead of 8.

This causes eval_loss to be computed as (loss(A,B,C,D) + loss(E,F,G,H) + loss(A,B,C,D)) / 2. Other metrics are also computed on A, B, C, D, E, F, G, H, A, B, C, D instead of only on A, B, C, D, E, F, G, H.

If the number of examples is not divisible by eval_batch_size, the first batch and part of the second batch are predicted twice. Say we have examples (n=8): A, B, C, D, E, F, G, H and eval_batch_size = 5. The current implementation of prediction_loop() will run batch (A,B,C,D,E), (F,G,H,A,B) and (C,D,E,F,G). This produces an output with a shape of 15 instead of 8. Again, the eval_loss and other metrics are computed incorrectly.

[This](https://colab.research.google.com/drive/1JH-269TcWWzowDngCmtpEwnmsFbVtL7K) is an example on an actual dataset (code is taken from the run_tf_glue.py example). Notice the assertion error when comparing the number of predicted results to the number of examples in the dataset.

I did realize that I made a complicated solution to the problem. The main change needed was to not call repeat(). I have committed the simpler solution.

Ok, from what I understand of the problem, what you have done is still not acceptable sorry, the build of a dataset must stay as it is because, `repeat` is very important and is mandatory when training on TPU, and `drop_remainder` must stay accessible through the training argument `dataloader_drop_last`, what if someone want it to be `False`?

If the problem is in the `prediction_loop` and goes one step to far, you can just stop the loop one step before by replacing:
```
if step == steps:
    break
```
by
```
if step == steps: - 1:
    break
```

This should work.

` if step == steps: - 1` only fixes the problem when `num_examples % batch_size`.

I am sorry if I missed this, but it seems like `prediction_loop` is only called during evaluation/prediction. I am not quite sure how it affects the training process. I did not change `get_train_tfdataset()`; `repeat()` is still in it. At the very least, `repeat()` should be removed in `get_test_tfdataset()` (i.e.during prediction), although I argue that it should be removed in 
`get_test_tfdataset()` too because the evaluation loss being reported is incorrect.


After giving a deeper look at the issue, I can see three things to fix:

1. Replace `if step == steps:` by `if step == steps: - 1:` line 348
2. Replace `metrics["eval_loss"] = self.eval_loss.result().numpy() / steps` by `metrics["eval_loss"] = self.eval_loss.result().numpy() / (steps - 1)` line 356
3. Move `self.eval_loss = tf.keras.metrics.Sum()` from the `prediction_loop` method inside the `__init__` method.

After having done those changes, run a training with the `--dataloader_drop_last` argument. Now you should not see the `0.0` loss value anymore.

The argument `--dataloader_drop_last` removes the last batch of the dataset. In order to be sure of that you can know the real size of your dataset by doing `(dataset.cardinality().numpy()// eval_batch_size) * eval_batch_size`. In the case of the MRPC dataset, `dataset.cardinality().numpy() == 408`, while the effective number of examples on which you will evaluate your model is `400`.

Thank you for fixing the 0.0 loss value!

Unfortunately, I think the first issue I mentioned still persists even with your first two fixes. If my math is correct, with your current solution, you use an `eval_batch_size` of 10 when evaluating the MRPC dataset, your model will be evaluated on 410 examples (the loss from the first two examples will be added twice to `eval_loss `). One way to check this is to log/print the shape of `preds`.

I am sorry if my explanations are confusing. I think my main point is that if I want to evaluate/predict on X examples, I expect it to evaluate/predict on exactly X examples, i.e. `preds.shape = (X, n_tags)` . I actually found this issue when I was using my trained model to predict 929 examples. I was getting 944 predictions, i.e. `predictions.shape = (944, n_tags)`. With your solution, I will be getting 936 examples.

On a related note, may I ask if there are unit tests for TFTrainer and if so, where these are located? I was only able to locate the Trainer tests.

Also, it seems like `drop_remainder` does not have an effect if you are using `repeat()`.

```python
import tensorflow as tf

dummy_data = tf.data.Dataset.range(8)
```

```python
test_batches = (
      dummy_data.repeat()
      .batch(3, drop_remainder=True)
      .prefetch(tf.data.experimental.AUTOTUNE)
  )

steps = 5
for step, batch in enumerate(test_batches):
  print(batch)
  if step == steps:
    break
```

```python
test_batches_nodrop = (
      dummy_data.repeat()
      .batch(3, drop_remainder=False)
      .prefetch(tf.data.experimental.AUTOTUNE)
  )

print('Batches when drop_remainder=False')
steps = 5
for step, batch in enumerate(test_batches_nodrop):
  print(batch)
  if step == steps:
    break

```

Both code blocks produce the same output:

```python
tf.Tensor([0 1 2], shape=(3,), dtype=int64)
tf.Tensor([3 4 5], shape=(3,), dtype=int64)
tf.Tensor([6 7 0], shape=(3,), dtype=int64)
tf.Tensor([1 2 3], shape=(3,), dtype=int64)
tf.Tensor([4 5 6], shape=(3,), dtype=int64)
tf.Tensor([7 0 1], shape=(3,), dtype=int64)
```

It does have effect on the behavior of `prediction_loop()` because it changes the `steps` to `steps-1`, but that is due to this line:

`approx = math.floor if self.args.dataloader_drop_last else math.ceil`

> I am sorry if my explanations are confusing. I think my main point is that if I want to evaluate/predict on X examples, I expect it to evaluate/predict on exactly X examples, i.e. preds.shape = (X, n_tags) .

Yes, but for this you have to use a compliant batch size, the requirements are to drop the last batch if its size is lower than the required batch size. This is the wanted and expected behavior. See my explanation on MRPC, with a batch size of 16 only 400 examples over 408 will be evaluated. If you want to evaluate over all the examples, you can use a batch size of 8.

> Also, it seems like drop_remainder does not have an effect if you are using repeat().

Yes, this is normal, as detailed in the documentation. This is one of the reason why we use the approx variable.


I see. Thank you for the explanation, and I am sorry for overlooking the comment about `drop_remainder` not having an effect. In this case, may I ask if it is ok to log the actual number of examples the model is evaluated on, e.g. adding something like: "Number of examples used for evaluation = 400"?

Does `predict()` require a similar behavior, e.g. is `repeat()` required in `get_test_tfdataset()`?  Unlike `evaluate()` it is never used in training. If I set `dataloader_drop_last=True` during training then perform prediction on unlabeled examples after, only 928 of my 929 examples are given a prediction.

> I see. Thank you for the explanation, and I am sorry for overlooking the comment about drop_remainder not having an effect.

No worries, that's ok :)

> In this case, may I ask if it is ok to log the actual number of examples the model is evaluated on, e.g. adding something like: "Number of examples used for evaluation = 400"?

Sure! That would be a good idea!!

> Does predict() require a similar behavior, e.g. is repeat() required in get_test_tfdataset()? Unlike evaluate() it is never used in training. If I set dataloader_drop_last=True during training then perform prediction on unlabeled examples after, only 928 of my 929 examples are given a prediction.

I agree, no need to use `repeat` for predict that uses the test dataset 👍 

I have pushed the discussed changes. A couple of final things:

> 2. Replace `metrics["eval_loss"] = self.eval_loss.result().numpy() / steps` by `metrics["eval_loss"] = self.eval_loss.result().numpy() / (steps - 1)` line 356

This does not need to be replaced. Since the loop terminates when `step== steps - 1`, it actually runs for n=steps times.

> 3. Move `self.eval_loss = tf.keras.metrics.Sum()` from the `prediction_loop` method inside the `__init__` method.

This works but I had to call `eval_loss.reset_states()` inside `prediction_loop()` so the sum from the previous calculations is not added.