Rework TF trainer
Author: jplu
Pull request: yes

This PR brings several improvements to the TensorFlow Trainer:

-  make it more intuitive with less multiple small functions.
- should also make the trainer compliant with TPU.
- Fix the TensorFlow version to at least 2.2 for the trainer
- increase speed of the dataset preprocessing

@Colanim Can you please test this PR with TPU to see if it looks ok.

Comments:
# [Codecov](https://codecov.io/gh/huggingface/transformers/pull/6038?src=pr&el=h1) Report
> Merging [#6038](https://codecov.io/gh/huggingface/transformers/pull/6038?src=pr&el=desc) into [master](https://codecov.io/gh/huggingface/transformers/commit/dafa296c952c08fca3686f1cf8f3a8f8eb116744&el=desc) will **decrease** coverage by `0.88%`.
> The diff coverage is `9.23%`.

[![Impacted file tree graph](https://codecov.io/gh/huggingface/transformers/pull/6038/graphs/tree.svg?width=650&height=150&src=pr&token=9qOlN6Hb1c)](https://codecov.io/gh/huggingface/transformers/pull/6038?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##           master    #6038      +/-   ##
==========================================
- Coverage   78.80%   77.92%   -0.89%     
==========================================
  Files         146      146              
  Lines       26325    26332       +7     
==========================================
- Hits        20746    20519     -227     
- Misses       5579     5813     +234     
```


| [Impacted Files](https://codecov.io/gh/huggingface/transformers/pull/6038?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [src/transformers/training\_args\_tf.py](https://codecov.io/gh/huggingface/transformers/pull/6038/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy90cmFpbmluZ19hcmdzX3RmLnB5) | `47.45% <0.00%> (ø)` | |
| [src/transformers/trainer\_tf.py](https://codecov.io/gh/huggingface/transformers/pull/6038/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy90cmFpbmVyX3RmLnB5) | `16.54% <8.66%> (+0.06%)` | :arrow_up: |
| [src/transformers/modeling\_tf\_utils.py](https://codecov.io/gh/huggingface/transformers/pull/6038/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ190Zl91dGlscy5weQ==) | `89.47% <100.00%> (+0.51%)` | :arrow_up: |
| [src/transformers/modeling\_tf\_mobilebert.py](https://codecov.io/gh/huggingface/transformers/pull/6038/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ190Zl9tb2JpbGViZXJ0LnB5) | `23.38% <0.00%> (-73.39%)` | :arrow_down: |
| [src/transformers/modeling\_tf\_gpt2.py](https://codecov.io/gh/huggingface/transformers/pull/6038/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ190Zl9ncHQyLnB5) | `65.42% <0.00%> (-29.91%)` | :arrow_down: |
| [src/transformers/generation\_tf\_utils.py](https://codecov.io/gh/huggingface/transformers/pull/6038/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9nZW5lcmF0aW9uX3RmX3V0aWxzLnB5) | `85.96% <0.00%> (-0.51%)` | :arrow_down: |
| [src/transformers/file\_utils.py](https://codecov.io/gh/huggingface/transformers/pull/6038/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9maWxlX3V0aWxzLnB5) | `82.48% <0.00%> (+0.28%)` | :arrow_up: |
| [src/transformers/modeling\_reformer.py](https://codecov.io/gh/huggingface/transformers/pull/6038/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ19yZWZvcm1lci5weQ==) | `94.97% <0.00%> (+4.41%)` | :arrow_up: |
| [src/transformers/modeling\_tf\_distilbert.py](https://codecov.io/gh/huggingface/transformers/pull/6038/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ190Zl9kaXN0aWxiZXJ0LnB5) | `98.79% <0.00%> (+33.89%)` | :arrow_up: |
| [src/transformers/modeling\_tf\_roberta.py](https://codecov.io/gh/huggingface/transformers/pull/6038/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ190Zl9yb2JlcnRhLnB5) | `93.36% <0.00%> (+49.37%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/huggingface/transformers/pull/6038?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/huggingface/transformers/pull/6038?src=pr&el=footer). Last update [dafa296...0e56209](https://codecov.io/gh/huggingface/transformers/pull/6038?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


@jplu I tried, but unfortunately I'm using a **TFRecord dataset**, which has an **unknown cardinality** ([doc link](https://www.tensorflow.org/api_docs/python/tf/data/experimental/cardinality)).

So I'm meeting the following error :

```
ValueError: The training dataset must have an asserted cardinality
```

Yes, you now have to use `assert_cardinality` on your dataset, as it is done in all the examples now.

I can rework it to make it more similar to the PT one, but I don't think they must be strictly the same has they don't really have the same way to work... In any case I can rethink it :)

About the issue I don't get what it means so I asked to elaborate more.

@jplu With my custom code I'm meeting the following error :

```
TypeError: in user code:

    /home/me/.venv/test/lib/python3.7/site-packages/transformers/trainer_tf.py:551 distributed_training_steps  *
        self.args.strategy.experimental_run_v2(apply_gradients, batch)
    /home/me/.venv/test/lib/python3.7/site-packages/transformers/trainer_tf.py:531 apply_gradients  *
        reduced_features = features[: self.args.train_batch_size / self.args.n_replicas]

    TypeError: unhashable type: 'slice'
```

Not sure if it's from my code or this PR though...

Humm this is weird... I never got this error myself. How did you come to this? in order to try to replicate it.

@LysandreJik @sgugger Did I properly address all your comments?