Fix a few countings (steps / epochs) in trainer_tf.py
Author: ydshieh
Pull request: yes

Mainly for @jplu , but @sgugger might have some comments, especially for issue 2.

## Description

This PR fixes a few countings (steps / epochs) in [trainer_tf.py](https://github.com/huggingface/transformers/blob/master/src/transformers/trainer_tf.py).

Since these issues are tight together, I decided to fix them in one PR.

BTW, I didn't find a `test_tf_trainer.py` similar to `test_trainer.py`, so no test added. I just relaunched the commands to make sure the fixed code give the desired outputs.

I put some comments in the code --> Some of them might be removed after the code is reviewed.

Here is the list of issues and their causes:

1. Potential `ZeroDivisionError` for lines 498 - 501 (similar to PR #7125  )

        epochs_trained = self.global_step // (self.num_train_examples // self.args.gradient_accumulation_steps)
        steps_trained_in_current_epoch = self.global_step % (
            self.num_train_examples // self.args.gradient_accumulation_steps
        )

2. The calculation of `epochs_trained` and `steps_trained_in_current_epoch` in 1. should be something like

        self.global_step // (self.num_train_examples // self.total_train_batch_size)

    instead, rather than using `self.args.gradient_accumulation_steps`. In pytorch trainer, we have (line 609)

        num_update_steps_per_epoch = len(train_dataloader) // self.args.gradient_accumulation_steps

   but there, `len(train_dataloader)` is the number of batches, not the number of examples.

3. Line 492, 498 and 499

        self.global_step = iterations.numpy()
        ....
       epochs_trained = self.global_step // (self.num_train_examples // self.args.gradient_accumulation_steps)
       steps_trained_in_current_epoch = self.global_step % (
           self.num_train_examples // self.args.gradient_accumulation_steps
       )

    count some training progresses before the checkpoint is restored (line 511)

         ckpt.restore(self.model.ckpt_manager.latest_checkpoint).expect_partial()

    These countings should be done after the optimizer is restored, otherewise they will be always zero at those lines.

4. Line 517

        epochs = 1 if self.args.max_steps > 0 else self.args.num_train_epochs

     make epochs to be `1` if `max_steps` is specified. This won't be a real bug, but when `max_steps` is larger than the optimizations steps per epochs, this is somehow confusing, especially when the log is shown.

5. Extra epoch issue:  Line 542

        for epoch_iter in range(epochs_trained, int(epochs + 1)):

    If `max_step` is not specified, and `num_train_epochs` is set to 1. Suppose in a first run, the model is trained for 1 step. Then when we resume the training, we will get `epochs_trained = 0` (see issue 2. above) and `epochs = 1`, there for the range becomes `range(0, 2)` which gives 2 epochs.

## Code showing the bugs

1. First download glue tasks

        python utils/download_glue_data.py --data_dir ./examples/text-classification/glue/ --tasks all

2. This code

    
        python3 run_tf_glue.py \
            --task_name wnli \
            --data_dir ./glue/WNLI \
            --model_name_or_path distilbert-base-uncased \
            --output_dir ./glue/WNLI/ \
            --max_seq_length  16 \
            --per_device_train_batch_size 16 \
            --gradient_accumulation_steps 4 \
            --max_steps 10 \
            --logging_steps 1 \
            --save_steps 5 \
            --seed 1 \
            --do_train \
            --do_eval \
            --do_predict \
            --overwrite_output_dir

    gives logs like

            {'loss': 0.69778687, 'learning_rate': 4.4999997e-05, 'epoch': 0.2, 'step': 1}
            {'loss': 0.6953752, 'learning_rate': 4e-05, 'epoch': 0.3, 'step': 2}

    where `'epoch': 0.2, 'step': 1` is not correct for `epoch`, because they are only 10 steps.

3. Remove the checkpoints in the output_dir. Launch 

        python3 run_tf_glue.py \
            --task_name wnli \
            --data_dir ./glue/WNLI \
            --model_name_or_path distilbert-base-uncased \
            --output_dir ./glue/WNLI/ \
            --max_seq_length  16 \
            --num_train_epochs 1 \
            --per_device_train_batch_size 16 \
            --gradient_accumulation_steps 4 \
            --max_steps 2 \
            --logging_steps 1 \
            --save_steps 1 \
            --seed 1 \
            --do_train \
            --do_eval \
            --do_predict \
            --overwrite_output_dir

    but stop it after 1 step. Then resume the training, we get logs like

        {'loss': 0.7005814, 'learning_rate': 0.0, 'epoch': 2.0, 'step': 1}

and 

        {'loss': 0.66841304, 'learning_rate': 0.0, 'epoch': -0.5, 'step': 2}
        Saving checkpoint for step 2 at ./glue/WNLI/checkpoint/ckpt-2
        {'loss': 0.7016809, 'learning_rate': 0.0, 'epoch': 0.5, 'step': 3}
        Saving checkpoint for step 3 at ./glue/WNLI/checkpoint/ckpt-3
        {'loss': 0.693136, 'learning_rate': 0.0, 'epoch': 1.0, 'step': 4}

The final steps became 4, and we got something like  `'epoch': -0.5`.

4. There are other codes to show the issues, but I hope the above two examples already make it clear that something is wrong in the current code. If you want, I can put more code snippet here.

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

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

```diff
@@            Coverage Diff             @@
##           master    #7175      +/-   ##
==========================================
+ Coverage   80.71%   80.99%   +0.27%     
==========================================
  Files         169      169              
  Lines       32293    32305      +12     
==========================================
+ Hits        26066    26165      +99     
+ Misses       6227     6140      -87     
```


| [Impacted Files](https://codecov.io/gh/huggingface/transformers/pull/7175?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [src/transformers/trainer\_tf.py](https://codecov.io/gh/huggingface/transformers/pull/7175/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy90cmFpbmVyX3RmLnB5) | `12.01% <0.00%> (-0.45%)` | :arrow_down: |
| [src/transformers/modeling\_tf\_longformer.py](https://codecov.io/gh/huggingface/transformers/pull/7175/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ190Zl9sb25nZm9ybWVyLnB5) | `16.37% <0.00%> (-82.31%)` | :arrow_down: |
| [src/transformers/modeling\_lxmert.py](https://codecov.io/gh/huggingface/transformers/pull/7175/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ19seG1lcnQucHk=) | `70.01% <0.00%> (-20.75%)` | :arrow_down: |
| [src/transformers/modeling\_roberta.py](https://codecov.io/gh/huggingface/transformers/pull/7175/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ19yb2JlcnRhLnB5) | `77.73% <0.00%> (-19.35%)` | :arrow_down: |
| [src/transformers/configuration\_bart.py](https://codecov.io/gh/huggingface/transformers/pull/7175/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9jb25maWd1cmF0aW9uX2JhcnQucHk=) | `90.00% <0.00%> (-4.00%)` | :arrow_down: |
| [src/transformers/data/data\_collator.py](https://codecov.io/gh/huggingface/transformers/pull/7175/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9kYXRhL2RhdGFfY29sbGF0b3IucHk=) | `92.83% <0.00%> (-0.72%)` | :arrow_down: |
| [src/transformers/pipelines.py](https://codecov.io/gh/huggingface/transformers/pull/7175/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9waXBlbGluZXMucHk=) | `80.75% <0.00%> (-0.25%)` | :arrow_down: |
| [src/transformers/tokenization\_utils\_base.py](https://codecov.io/gh/huggingface/transformers/pull/7175/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy90b2tlbml6YXRpb25fdXRpbHNfYmFzZS5weQ==) | `94.04% <0.00%> (+0.27%)` | :arrow_up: |
| [src/transformers/modeling\_bart.py](https://codecov.io/gh/huggingface/transformers/pull/7175/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ19iYXJ0LnB5) | `94.27% <0.00%> (+0.33%)` | :arrow_up: |
| [...rc/transformers/data/datasets/language\_modeling.py](https://codecov.io/gh/huggingface/transformers/pull/7175/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9kYXRhL2RhdGFzZXRzL2xhbmd1YWdlX21vZGVsaW5nLnB5) | `94.11% <0.00%> (+1.17%)` | :arrow_up: |
| ... and [6 more](https://codecov.io/gh/huggingface/transformers/pull/7175/diff?src=pr&el=tree-more) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/huggingface/transformers/pull/7175?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/7175?src=pr&el=footer). Last update [df16506...b72f0d9](https://codecov.io/gh/huggingface/transformers/pull/7175?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


Nice catch!! LGTM.

@LysandreJik I fully agree on adding tests for the trainer. I put this immediately on top of my todo list.

> 
> 
> @LysandreJik I fully agree on adding tests for the trainer. I put this immediately on top of my todo list.

@LysandreJik @jplu If you agree, may I help to build tests for the TFTrainer? If so, probably I need to discuss with @jplu how to proceed though.

@chiapas I will definitely be glad to have your help on this! Can you send me an email (it is displayed on my Github profile) please, and I will let you know how we can proceed.

If you could copy and adapt the tests in trainer (so that they run quickly) on a regression problem, it would be awesome!

> 
> 
> If you could copy and adapt the tests in trainer (so that they run quickly) on a regression problem, it would be awesome!

We are going to build the tests for TFTrainer soon. But the HF team can decide if to merge this PR now or after the tests are built.

LGTM for merge, what do you think @sgugger?

Yes, good to merge. Thanks @chiapas!