Repository: Lightning-AI/pytorch-lightning

Issue:
Unwanted accumulate_grad_batches behavior
## 🐛 Bug

When using the flag `accumulate_grad_batches` for the trainer, if an action is to be performed at the last mini-batch it isn't done.

### To Reproduce

Steps to reproduce the behavior:

1. In `on_after_backward`, define some logic if we are at the last mini-batch
``` 
if self.__nbbatch -1 <= self.__batchidx:
 some_param.grad += gradient_penalty
```
2. Run with and without the flag
3. If running with the flag, the gradient penalty has not effect (the optimizer probably didn't take a step for the last mini-batch)

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->


#### Code sample
See above. 
### Expected behavior

I manually compute a gradient penalty that is applied only at the last mini-batch of an epoch. Using the flag shouldn't break this behavior.


<!-- A clear and concise description of what you expected to happen. -->

### Environment

PyTorch version: 1.6.0.dev20200403
Is debug build: No
CUDA used to build PyTorch: 10.1

OS: CentOS Linux release 7.7.1908 (Core)
GCC version: (Homebrew GCC 5.5.0_7) 5.5.0
CMake version: version 3.13.0

Python version: 3.7
Is CUDA available: Yes
CUDA runtime version: 10.2.89
GPU models and configuration: 
GPU 0: Tesla P100-SXM2-16GB
GPU 1: Tesla P100-SXM2-16GB
GPU 2: Tesla P100-SXM2-16GB
GPU 3: Tesla P100-SXM2-16GB

Nvidia driver version: 440.33.01
cuDNN version: Could not collect

Versions of relevant libraries:
[pip3] numpy==1.18.1
[conda] blas                      1.0                         mkl  
[conda] kmeans-pytorch            0.2                      pypi_0    pypi
[conda] mkl                       2019.4                      243  
[conda] mkl-service               2.3.0            py37he904b0f_0  
[conda] mkl_fft                   1.0.15           py37ha843d7b_0  
[conda] mkl_random                1.1.0            py37hd6b4f25_0  
[conda] pytorch                   1.6.0.dev20200403 py3.7_cuda10.1.243_cudnn7.6.3_0    pytorch-nightly
[conda] pytorch-lightning         0.7.3                    pypi_0    pypi
[conda] pytorch-memlab            0.0.4                    pypi_0    pypi
[conda] pytorch-pcen              0.0.1                    pypi_0    pypi
[conda] torchvision               0.6.0.dev20200403      py37_cu101    pytorch-nightly

### Additional context

It isn't a bug per se, but it should at least be a documented behavior, ideally controllable with a flag.



Discussion:
This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.

@sebastienwood thanks for bringing this up! we're looking at it for next sprint

Resolving patch:
diff --git a/pytorch_lightning/callbacks/gradient_accumulation_scheduler.py b/pytorch_lightning/callbacks/gradient_accumulation_scheduler.py
--- a/pytorch_lightning/callbacks/gradient_accumulation_scheduler.py
+++ b/pytorch_lightning/callbacks/gradient_accumulation_scheduler.py
@@ -3,6 +3,7 @@
 ====================
 
 Change gradient accumulation factor according to scheduling.
+Trainer also calls ``optimizer.step()`` for the last indivisible step number.
 
 """
 
diff --git a/pytorch_lightning/trainer/__init__.py b/pytorch_lightning/trainer/__init__.py
--- a/pytorch_lightning/trainer/__init__.py
+++ b/pytorch_lightning/trainer/__init__.py
@@ -168,6 +168,7 @@ def forward(self, x):
 accumulate_grad_batches
 ^^^^^^^^^^^^^^^^^^^^^^^
 Accumulates grads every k batches or as set up in the dict.
+Trainer also calls ``optimizer.step()`` for the last indivisible step number.
 
 .. testcode::
 
diff --git a/pytorch_lightning/trainer/training_loop.py b/pytorch_lightning/trainer/training_loop.py
--- a/pytorch_lightning/trainer/training_loop.py
+++ b/pytorch_lightning/trainer/training_loop.py
@@ -590,7 +590,8 @@ def check_checkpoint_callback(self, should_check_val):
             [c.on_validation_end(self, self.get_model()) for c in checkpoint_callbacks]
 
     def update_train_loop_lr_schedulers(self):
-        if (self.batch_idx + 1) % self.accumulate_grad_batches == 0:
+        if ((self.batch_idx + 1) % self.accumulate_grad_batches == 0
+                or (self.batch_idx + 1) == self.num_training_batches):
             # update lr
             self.update_learning_rates(interval='step')
 
@@ -737,7 +738,8 @@ def sync_horovod(self):
 
     def increment_accumulated_grad_global_step(self):
         # progress global step according to grads progress
-        if (self.batch_idx + 1) % self.accumulate_grad_batches == 0:
+        if ((self.batch_idx + 1) % self.accumulate_grad_batches == 0
+                or (self.batch_idx + 1) == self.num_training_batches):
             self.global_step += 1
         self.total_batch_idx += 1
 
@@ -878,7 +880,8 @@ def run_training_batch(self, batch, batch_idx):
                 # BACKWARD PASS
                 # ------------------------------
                 # gradient update with accumulated gradients
-                if (self.batch_idx + 1) % self.accumulate_grad_batches == 0:
+                if ((self.batch_idx + 1) % self.accumulate_grad_batches == 0
+                        or (self.batch_idx + 1) == self.num_training_batches):
 
                     # backward
                     grad_norm_dic = self.run_batch_backward_pass(split_batch, batch_idx, opt_idx, optimizer)


Test patch:
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
@@ -5,6 +5,7 @@
 import sys
 import types
 from argparse import Namespace
+from copy import deepcopy
 from pathlib import Path
 from unittest.mock import patch
 
@@ -196,7 +197,7 @@ def test_gradient_accumulation_scheduling(tmpdir, schedule, expected):
 
     trainer = Trainer(
         accumulate_grad_batches=schedule,
-        limit_train_batches=0.8,
+        limit_train_batches=0.7,  # not to be divisible by accumulate_grad_batches on purpose
         limit_val_batches=0.8,
         max_epochs=4,
         default_root_dir=tmpdir,
@@ -216,8 +217,15 @@ def _optimizer_step(epoch, batch_idx, optimizer, optimizer_idx,
                     # use this opportunity to test once
                     assert trainer.accumulate_grad_batches == expected[0]
 
-                assert batch_idx == model.prev_called_batch_idx
-                model.prev_called_batch_idx += expected[0]
+                # separate check for last batch with accumulate 1 step
+                if expected[0] == 1 and (batch_idx + 1) == trainer.num_training_batches:
+                    assert batch_idx == model.prev_called_batch_idx
+                elif (batch_idx + 1) == trainer.num_training_batches:
+                    # prev_called_batch_idx - schedule + modulus remainder
+                    assert batch_idx == (model.prev_called_batch_idx - expected[0] + (batch_idx + 1) % expected[0])
+                else:
+                    assert batch_idx == model.prev_called_batch_idx
+                    model.prev_called_batch_idx += expected[0]
 
             elif 1 <= epoch <= 2:
                 # reset counter when starting epoch
@@ -227,8 +235,12 @@ def _optimizer_step(epoch, batch_idx, optimizer, optimizer_idx,
                     # use this opportunity to test once
                     assert trainer.accumulate_grad_batches == expected[1]
 
-                assert batch_idx == model.prev_called_batch_idx
-                model.prev_called_batch_idx += expected[1]
+                if trainer.num_training_batches == batch_idx + 1:
+                    # prev_called_batch_idx - schedule + modulus remainder
+                    assert batch_idx == (model.prev_called_batch_idx - expected[1] + (batch_idx + 1) % expected[1])
+                else:
+                    assert batch_idx == model.prev_called_batch_idx
+                    model.prev_called_batch_idx += expected[1]
 
             else:
                 if batch_idx == expected[2] - 1:
@@ -237,8 +249,12 @@ def _optimizer_step(epoch, batch_idx, optimizer, optimizer_idx,
                     # use this opportunity to test once
                     assert trainer.accumulate_grad_batches == expected[2]
 
-                assert batch_idx == model.prev_called_batch_idx
-                model.prev_called_batch_idx += expected[2]
+                if (batch_idx + 1) == trainer.num_training_batches:
+                    # prev_called_batch_idx - schedule + modulus remainder
+                    assert batch_idx == (model.prev_called_batch_idx - expected[2] + (batch_idx + 1) % expected[2])
+                else:
+                    assert batch_idx == model.prev_called_batch_idx
+                    model.prev_called_batch_idx += expected[2]
 
         optimizer.step()
 
@@ -252,6 +268,48 @@ def _optimizer_step(epoch, batch_idx, optimizer, optimizer_idx,
     trainer.fit(model)
 
 
+@pytest.mark.parametrize(
+    ['accumulate_grad_batches', 'limit_train_batches'],
+    [
+        pytest.param({1: 2, 3: 4}, 1.0),
+        pytest.param({1: 2, 3: 4}, 0.5),  # not to be divisible by accumulate_grad_batches on purpose
+        pytest.param(3, 1.0),
+        pytest.param(3, 0.8),  # not to be divisible by accumulate_grad_batches on purpose
+        pytest.param(4, 1.0),
+        pytest.param(4, 0.7),  # not to be divisible by accumulate_grad_batches on purpose
+    ],
+)
+def test_gradient_accumulation_scheduling_last_batch(tmpdir, accumulate_grad_batches, limit_train_batches):
+    """ Verify optimizer.step() applied to last batch while grad accumulation """
+
+    class CurrentModel(EvalModelTemplate):
+        def on_after_backward(self):
+            self.loss_backward = deepcopy(self.state_dict())
+
+        def on_before_zero_grad(self, optimizer):
+            self.opt_step = self.state_dict()
+
+        def on_train_batch_end(self, batch, batch_idx, dataloader_idx):
+            _exclude_keys = ['num_batches_tracked', 'running_mean', 'running_var']
+
+            if (batch_idx + 1) == self.trainer.num_training_batches:
+                for key in self.loss_backward.keys():
+                    # exclude the check for batch_norm parameters
+                    if not any([k in key for k in _exclude_keys]):
+                        assert not torch.equal(self.loss_backward[key], self.opt_step[key])
+
+    model = CurrentModel()
+
+    trainer = Trainer(
+        accumulate_grad_batches=accumulate_grad_batches,
+        max_epochs=4,
+        limit_train_batches=limit_train_batches,
+        default_root_dir=tmpdir
+    )
+
+    trainer.fit(model)
+
+
 def test_loading_meta_tags(tmpdir):
     """ test for backward compatibility to meta_tags.csv """
     tutils.reset_seed()