Update docs for GradientAccumulationScheduler (#9891)

* update docs and add tests

* update docs and add tests

* Update pytorch_lightning/callbacks/gradient_accumulation_scheduler.py

Co-authored-by: Rohit Gupta <[email protected]>

* Apply suggestions from code review

Co-authored-by: Adrian Wälchli <[email protected]>

Co-authored-by: Adrian Wälchli <[email protected]>

Author: rohitgr7

pytorch_lightning/callbacks/gradient_accumulation_scheduler.py

@@ -23,6 +23,7 @@
 from typing import Dict
 
 from pytorch_lightning.callbacks.base import Callback
+from pytorch_lightning.utilities.exceptions import MisconfigurationException
 
 
 class GradientAccumulationScheduler(Callback):
@@ -32,6 +33,14 @@ class GradientAccumulationScheduler(Callback):
     Args:
         scheduling: scheduling in format {epoch: accumulation_factor}
 
+    Note:
+        The argument scheduling is a dictionary. Each key represent an epoch and
+        its associated accumulation factor value.
+        Warning: Epoch are zero-indexed c.f it means if you want to change
+        the accumulation factor after 4 epochs, set ``Trainer(accumulate_grad_batches={4: factor})``
+        or ``GradientAccumulationScheduler(scheduling={4: factor})``.
+        For more info check the example below.
+
     Raises:
         TypeError:
             If ``scheduling`` is an empty ``dict``,
@@ -44,12 +53,13 @@ class GradientAccumulationScheduler(Callback):
         >>> from pytorch_lightning import Trainer
         >>> from pytorch_lightning.callbacks import GradientAccumulationScheduler
 
-        # at epoch 5 start accumulating every 2 batches
-        >>> accumulator = GradientAccumulationScheduler(scheduling={5: 2})
+        # from epoch 5, it starts accumulating every 2 batches. Here we have 4 instead of 5
+        because epoch (key) should be zero-indexed.
+        >>> accumulator = GradientAccumulationScheduler(scheduling={4: 2})
         >>> trainer = Trainer(callbacks=[accumulator])
 
         # alternatively, pass the scheduling dict directly to the Trainer
-        >>> trainer = Trainer(accumulate_grad_batches={5: 2})
+        >>> trainer = Trainer(accumulate_grad_batches={4: 2})
     """
 
     def __init__(self, scheduling: Dict[int, int]):
@@ -58,9 +68,15 @@ def __init__(self, scheduling: Dict[int, int]):
         if not scheduling:  # empty dict error
             raise TypeError("Empty dict cannot be interpreted correct")
 
-        for key in scheduling:
-            if not isinstance(key, int) or not isinstance(scheduling[key], int):
-                raise TypeError("All epoches and accumulation factor must be integers")
+        if any(not isinstance(key, int) or key < 0 for key in scheduling):
+            raise MisconfigurationException(
+                f"Epoch should be an int greater than or equal to 0. Got {list(scheduling.keys())}."
+            )
+
+        if any(not isinstance(value, int) or value < 1 for value in scheduling.values()):
+            raise MisconfigurationException(
+                f"Accumulation factor should be an int greater than 0. Got {list(scheduling.values())}."
+            )
 
         minimal_epoch = min(scheduling.keys())
         if minimal_epoch < 0:

tests/callbacks/test_gradient_accumulation_scheduler.py

@@ -0,0 +1,108 @@
+# Copyright The PyTorch Lightning team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import math
+from unittest.mock import patch
+
+import pytest
+
+from pytorch_lightning import Trainer
+from pytorch_lightning.callbacks import GradientAccumulationScheduler
+from pytorch_lightning.utilities.exceptions import MisconfigurationException
+from tests.helpers import BoringModel
+
+
+@pytest.mark.parametrize("accumulate_grad_batches", (1, 2, 3))
+def test_trainer_accumulate_grad_batches_zero_grad(tmpdir, accumulate_grad_batches):
+    with patch("torch.optim.SGD.zero_grad") as sgd_zero_grad:
+        model = BoringModel()
+        trainer = Trainer(
+            default_root_dir=tmpdir,
+            limit_train_batches=20,
+            limit_val_batches=1,
+            max_epochs=1,
+            weights_summary=None,
+            accumulate_grad_batches=accumulate_grad_batches,
+        )
+        assert trainer.accumulate_grad_batches == accumulate_grad_batches
+        trainer.fit(model)
+
+        assert sum(isinstance(cb, GradientAccumulationScheduler) for cb in trainer.callbacks) == 1
+        assert sgd_zero_grad.call_count == math.ceil(trainer.limit_train_batches / accumulate_grad_batches)
+
+
+@pytest.mark.parametrize(
+    ["accumulate_grad_batches", "expected_call_count"],
+    [
+        ({1: 2, 3: 4}, 10 + 5 + 5 + 3),
+        ({0: 2, 2: 1}, 5 + 5 + 10 + 10),
+    ],
+)
+def test_trainer_accumulate_grad_batches_dict_zero_grad(tmpdir, accumulate_grad_batches, expected_call_count):
+    with patch("torch.optim.SGD.zero_grad") as sgd_zero_grad:
+        model = BoringModel()
+        trainer = Trainer(
+            default_root_dir=tmpdir,
+            limit_train_batches=10,
+            limit_val_batches=1,
+            max_epochs=4,
+            weights_summary=None,
+            accumulate_grad_batches=accumulate_grad_batches,
+        )
+        assert trainer.accumulate_grad_batches == accumulate_grad_batches.get(0, 1)
+        trainer.fit(model)
+
+        assert sum(isinstance(cb, GradientAccumulationScheduler) for cb in trainer.callbacks) == 1
+        assert sgd_zero_grad.call_count == expected_call_count
+
+
+def test_trainer_accumulate_grad_batches_with_callback(tmpdir):
+    with patch("torch.optim.SGD.zero_grad") as sgd_zero_grad:
+        model = BoringModel()
+        trainer = Trainer(
+            default_root_dir=tmpdir,
+            limit_train_batches=10,
+            limit_val_batches=1,
+            max_epochs=4,
+            weights_summary=None,
+            callbacks=[GradientAccumulationScheduler({1: 2, 3: 4})],
+        )
+        assert trainer.accumulate_grad_batches == 1
+        trainer.fit(model)
+
+        assert sum(isinstance(cb, GradientAccumulationScheduler) for cb in trainer.callbacks) == 1
+        assert sgd_zero_grad.call_count == 10 + 5 + 5 + 3
+
+
+@pytest.mark.parametrize(
+    "scheduling",
+    [
+        {1: 2, -3: 4},
+        {0: 2, "2": 1},
+    ],
+)
+def test_invalid_keys_for_grad_accum_scheduler(scheduling):
+    with pytest.raises(MisconfigurationException, match="Epoch should be an int"):
+        _ = GradientAccumulationScheduler(scheduling=scheduling)
+
+
+@pytest.mark.parametrize(
+    "scheduling",
+    [
+        {1: 0, 3: 4},
+        {0: 2, 2: "2"},
+    ],
+)
+def test_invalid_values_for_grad_accum_scheduler(scheduling):
+    with pytest.raises(MisconfigurationException, match="Accumulation factor should be an int"):
+        _ = GradientAccumulationScheduler(scheduling=scheduling)

tests/trainer/test_trainer.py

@@ -174,68 +174,6 @@ def test_strict_model_load(monkeypatch, tmpdir, tmpdir_server, url_ckpt):
     assert not failed, "Model should be loaded due to strict=False."
 
 
-@pytest.mark.parametrize("accumulate_grad_batches", (1, 2, 3))
-def test_trainer_accumulate_grad_batches_zero_grad(tmpdir, accumulate_grad_batches):
-    with patch("torch.optim.SGD.zero_grad") as sgd_zero_grad:
-        model = BoringModel()
-        trainer = Trainer(
-            default_root_dir=tmpdir,
-            limit_train_batches=20,
-            limit_val_batches=1,
-            max_epochs=1,
-            weights_summary=None,
-            accumulate_grad_batches=accumulate_grad_batches,
-        )
-        assert trainer.accumulate_grad_batches == accumulate_grad_batches
-        trainer.fit(model)
-
-        assert sum(isinstance(cb, GradientAccumulationScheduler) for cb in trainer.callbacks) == 1
-        assert sgd_zero_grad.call_count == math.ceil(trainer.limit_train_batches / accumulate_grad_batches)
-
-
-@pytest.mark.parametrize(
-    ["accumulate_grad_batches", "expected_call_count"],
-    [
-        ({1: 2, 3: 4}, 10 + 5 + 5 + 3),
-        ({0: 2, 2: 1}, 5 + 5 + 10 + 10),
-    ],
-)
-def test_trainer_accumulate_grad_batches_dict_zero_grad(tmpdir, accumulate_grad_batches, expected_call_count):
-    with patch("torch.optim.SGD.zero_grad") as sgd_zero_grad:
-        model = BoringModel()
-        trainer = Trainer(
-            default_root_dir=tmpdir,
-            limit_train_batches=10,
-            limit_val_batches=1,
-            max_epochs=4,
-            weights_summary=None,
-            accumulate_grad_batches=accumulate_grad_batches,
-        )
-        assert trainer.accumulate_grad_batches == accumulate_grad_batches.get(0, 1)
-        trainer.fit(model)
-
-        assert sum(isinstance(cb, GradientAccumulationScheduler) for cb in trainer.callbacks) == 1
-        assert sgd_zero_grad.call_count == expected_call_count
-
-
-def test_trainer_accumulate_grad_batches_with_callback(tmpdir):
-    with patch("torch.optim.SGD.zero_grad") as sgd_zero_grad:
-        model = BoringModel()
-        trainer = Trainer(
-            default_root_dir=tmpdir,
-            limit_train_batches=10,
-            limit_val_batches=1,
-            max_epochs=4,
-            weights_summary=None,
-            callbacks=[GradientAccumulationScheduler({1: 2, 3: 4})],
-        )
-        assert trainer.accumulate_grad_batches == 1
-        trainer.fit(model)
-
-        assert sum(isinstance(cb, GradientAccumulationScheduler) for cb in trainer.callbacks) == 1
-        assert sgd_zero_grad.call_count == 10 + 5 + 5 + 3
-
-
 def test_trainer_accumulate_grad_batches_incorrect_value(tmpdir):
     with pytest.raises(MisconfigurationException, match=".*should be an int or a dict.*"):
         Trainer(default_root_dir=tmpdir, accumulate_grad_batches=(2, 5))