Repository: Lightning-AI/pytorch-lightning

Issue:
For versions >0.8.2 learning rate is zero for last epoch (potentially a logging bug)
## 🐛 Bug

Version 0.8.2 and above changed the behavior of either my learning rate scheduler or the `WandbLogger` logger. I am using a linear warmup and decay scheduler. However, the learning rate graph produced by the `LearningRateLogger` is as shown below ever since version 0.8.2:

![image1](https://user-images.githubusercontent.com/11785397/86428929-fc898a80-bcbb-11ea-85b1-2a42f4c54cfd.png)

The period where the learning rate is zero corresponds to the last epoch of training as you can see below:

![image2](https://user-images.githubusercontent.com/11785397/86428981-22af2a80-bcbc-11ea-90ff-bf4418dfe090.png)

This graph raises another issue. The first epoch appears to take twice as many steps as the second and third epoch. I specified `max_epochs=3`. During training, each epoch takes the same amount of time, so this seems like a logging issue.

Note that the above graphs are for a model that had its training stopped early. So the last epoch is slightly shorter than the second to last. This is not the issue.

Both of these issues (the 0 learning rate and the twice-as-long epoch) do not exist in version 0.8.1, and both graphs look as they should. 

These issues could be caused by the logger or they might actually occur and be logged correctly. I have looked through the [changelog](https://github.com/PyTorchLightning/pytorch-lightning/releases) and I am guessing that these bugs are caused by "Changed epoch indexing from 0 instead of 1" (#2289). I also may be relying on the fact that epoch indexing started at 1 somewhere in my code, but I do not believe this to be the case.

### To Reproduce

Reproducing this problem may be difficult since I can't provide the script and data I used. I used the `WandbLogger` logger and `LearningRateLogger` callback. I trained with 1400 warmup steps and `accumulate_grad_batches` set to 2.

I can provide additional code samples or information that you may need.

#### Code sample
```python
def lr_lambda_func(current_step, num_warmup_steps, num_training_steps):
    if current_step < num_warmup_steps:
        return float(current_step) / float(max(1, num_warmup_steps))
    return max(
        0.0,
        float(num_training_steps - current_step)
        / float(max(1, num_training_steps - num_warmup_steps)),
    )

t_total = int(len(self.train_dataloader_object) * self.hparams.max_epochs // self.hparams.accumulate_grad_batches)

lr_lambda = partial(
    lr_lambda_func,
    num_warmup_steps=self.hparams.warmup_steps
    * self.hparams.accumulate_grad_batches,
    num_training_steps=t_total,
)

scheduler = LambdaLR(optimizer, lr_lambda, -1)
scheduler_dict = {"scheduler": scheduler, "interval": "step"}
return ([optimizer], [scheduler_dict])
```

### Expected behavior

The learning rate should warmup and decay in versions greater than 0.8.2 the same way it does in versions less than 0.8.2. Each epoch should be the same number of steps.

The below graphs highlight the expected behavior. They are from a different model so they are not directly comparable, but their shape is as expected since they were captured from a model trained with `pytorch_lightning` version 0.8.1.

![image3](https://user-images.githubusercontent.com/11785397/86430600-b5ea5f00-bcc0-11ea-8c07-b78670f1ef7d.png)

![image4](https://user-images.githubusercontent.com/11785397/86430629-c8649880-bcc0-11ea-8aca-85cdde47a0d2.png)

### Environment

* CUDA:
	- GPU:
		- Tesla P100-PCIE-16GB
	- available:         True
	- version:           10.1
* Packages:
	- numpy:             1.18.5
	- pyTorch_debug:     False
	- pyTorch_version:   1.5.1+cu101
	- pytorch-lightning: 0.8.4
	- tensorboard:       2.2.2
	- tqdm:              4.41.1
* System:
	- OS:                Linux
	- architecture:
		- 64bit
		- 
	- processor:         x86_64
	- python:            3.6.9
	- version:           #1 SMP Wed Feb 19 05:26:34 PST 2020




Discussion:
Hi! thanks for your contribution!, great first issue!
it would be good to know whether this can be observed with the other loggers as well. Could you run your example also with TensorboardLogger?
Hey! I believe problem lies in `configure_accumulated_gradients()` when `accumulate_grad_batches` is integer, scheduler is set to use it from `current_epoch=1`, but Trainer starts from `current_epoch=0`, so `trainer.accumulate_grad_batches = self.scheduling.get(self.epochs[i])` sets `accumulate_grad_batches` to default value (`1`) for this epoch.
Update: When `accumulate_grad_batches` is an integer, Scheduler gets `{1: accumulate_grad_batches}` as input, and then `scheduling.update({0: 1})` inserts "default" `1` for first epoch.
@HHousen You could do workaround and set `accumulate_grad_batches={0: <your value>}` in pl.Trainer (I did so), but you might have problems with restoring from checkpoint, as 
```
n_accum = 1 if self.accumulate_grad_batches is None else self.accumulate_grad_batches
expected_steps = self.num_training_batches / n_accum
```
in `restore_training_state` will try to use dict in division.
@szymonzareba Yep, setting `args.accumulate_grad_batches` to `{0: 2}` fixed this problem (I create my `pl.Trainer` like so: `trainer = Trainer.from_argparse_args(args)`). Both the learning rate and epoch graphs are now correct. It seems like your reasoning is correct.
@HHousen mind send a PR? 
Sure

Resolving patch:
diff --git a/pytorch_lightning/trainer/training_tricks.py b/pytorch_lightning/trainer/training_tricks.py
--- a/pytorch_lightning/trainer/training_tricks.py
+++ b/pytorch_lightning/trainer/training_tricks.py
@@ -99,7 +99,7 @@ def configure_accumulated_gradients(self, accumulate_grad_batches):
         if isinstance(accumulate_grad_batches, dict):
             self.accumulation_scheduler = GradientAccumulationScheduler(accumulate_grad_batches)
         elif isinstance(accumulate_grad_batches, int):
-            schedule = {1: accumulate_grad_batches}
+            schedule = {0: accumulate_grad_batches}
             self.accumulation_scheduler = GradientAccumulationScheduler(schedule)
         else:
             raise TypeError("Gradient accumulation supports only int and dict types")


Test patch:
diff --git a/tests/trainer/test_lr_finder.py b/tests/trainer/test_lr_finder.py
--- a/tests/trainer/test_lr_finder.py
+++ b/tests/trainer/test_lr_finder.py
@@ -154,7 +154,7 @@ def test_accumulation_and_early_stopping(tmpdir):
         'Learning rate was not altered after running learning rate finder'
     assert len(lrfinder.results['lr']) == 100, \
         'Early stopping for learning rate finder did not work'
-    assert lrfinder._total_batch_idx == 190, \
+    assert lrfinder._total_batch_idx == 100 * 2, \
         'Accumulation parameter did not work'
 
 
diff --git a/tests/trainer/test_trainer.py b/tests/trainer/test_trainer.py
--- a/tests/trainer/test_trainer.py
+++ b/tests/trainer/test_trainer.py
@@ -103,7 +103,15 @@ def test_no_val_end_module(monkeypatch, tmpdir, tmpdir_server, url_ckpt):
     model_2.eval()
 
 
-def test_gradient_accumulation_scheduling(tmpdir):
+@pytest.mark.parametrize(
+    ['schedule', 'expected'],
+    [
+        pytest.param({1: 2, 3: 4}, [1, 2, 4]),
+        pytest.param(3, [3, 3, 3]),
+        pytest.param(4, [4, 4, 4])
+    ]
+)
+def test_gradient_accumulation_scheduling(tmpdir, schedule, expected):
     """
     Test grad accumulation by the freq of optimizer updates
     """
@@ -123,59 +131,59 @@ def test_gradient_accumulation_scheduling(tmpdir):
     with pytest.raises(TypeError):
         assert Trainer(accumulate_grad_batches={1: 2.5, 3: 5})
 
+    model = EvalModelTemplate()
+
+    trainer = Trainer(accumulate_grad_batches=schedule,
+                      limit_train_batches=0.8,
+                      limit_val_batches=0.8,
+                      max_epochs=4,
+                      default_root_dir=tmpdir)
+
     # test optimizer call freq matches scheduler
-    def _optimizer_step(self, epoch, batch_idx, optimizer,
-                        optimizer_idx, second_order_closure=None):
+    def _optimizer_step(epoch, batch_idx, optimizer, optimizer_idx,
+                        second_order_closure=None, on_tpu=False,
+                        using_native_amp=False, using_lbfgs=False):
         # only test the first 12 batches in epoch
         if batch_idx < 12:
             if epoch == 0:
                 # reset counter when starting epoch
-                if batch_idx == 0:
-                    self.prev_called_batch_idx = 0
+                if batch_idx == expected[0] - 1:
+                    model.prev_called_batch_idx = expected[0] - 1
 
                     # use this opportunity to test once
-                    assert self.trainer.accumulate_grad_batches == 1
+                    assert trainer.accumulate_grad_batches == expected[0]
 
-                assert batch_idx == self.prev_called_batch_idx
-                self.prev_called_batch_idx += 1
+                assert batch_idx == model.prev_called_batch_idx
+                model.prev_called_batch_idx += expected[0]
 
             elif 1 <= epoch <= 2:
                 # reset counter when starting epoch
-                if batch_idx == 1:
-                    self.prev_called_batch_idx = 1
+                if batch_idx == expected[1] - 1:
+                    model.prev_called_batch_idx = expected[1] - 1
 
                     # use this opportunity to test once
-                    assert self.trainer.accumulate_grad_batches == 2
+                    assert trainer.accumulate_grad_batches == expected[1]
 
-                assert batch_idx == self.prev_called_batch_idx
-                self.prev_called_batch_idx += 2
+                assert batch_idx == model.prev_called_batch_idx
+                model.prev_called_batch_idx += expected[1]
 
             else:
-                if batch_idx == 3:
-                    self.prev_called_batch_idx = 3
+                if batch_idx == expected[2] - 1:
+                    model.prev_called_batch_idx = expected[2] - 1
 
                     # use this opportunity to test once
-                    assert self.trainer.accumulate_grad_batches == 4
+                    assert trainer.accumulate_grad_batches == expected[2]
 
-                assert batch_idx == self.prev_called_batch_idx
-                self.prev_called_batch_idx += 3
+                assert batch_idx == model.prev_called_batch_idx
+                model.prev_called_batch_idx += expected[2]
 
         optimizer.step()
 
         # clear gradients
         optimizer.zero_grad()
 
-    model = EvalModelTemplate()
-    schedule = {1: 2, 3: 4}
-
-    trainer = Trainer(accumulate_grad_batches=schedule,
-                      limit_train_batches=0.1,
-                      limit_val_batches=0.1,
-                      max_epochs=2,
-                      default_root_dir=tmpdir)
-
     # for the test
-    trainer.optimizer_step = _optimizer_step
+    model.optimizer_step = _optimizer_step
     model.prev_called_batch_idx = 0
 
     trainer.fit(model)