Add unit tests for pl.utilities.grads (#9765)

Co-authored-by: Carlos Mocholí <[email protected]>

Author: awaelchli

CHANGELOG.md

@@ -297,12 +297,17 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Update the logic to check for accumulation steps with deepspeed ([#9826](https://github.com/PyTorchLightning/pytorch-lightning/pull/9826))
 
 
+- `pytorch_lightning.utilities.grads.grad_norm` now raises an exception if parameter `norm_type <= 0` ([#9765](https://github.com/PyTorchLightning/pytorch-lightning/pull/9765))
+
+
+
 - Updated error message for interactive incompatible plugins ([#9896](https://github.com/PyTorchLightning/pytorch-lightning/pull/9896))
 
 
 - Updated several places in the loops and trainer to access `training_type_plugin` directly instead of `accelerator` ([#9901](https://github.com/PyTorchLightning/pytorch-lightning/pull/9901))
 
 
+
 ### Deprecated
 
 - Deprecated trainer argument `terminate_on_nan` in favour of `detect_anomaly`([#9175](https://github.com/PyTorchLightning/pytorch-lightning/pull/9175))

pytorch_lightning/trainer/connectors/training_trick_connector.py

@@ -49,9 +49,11 @@ def on_trainer_init(
             )
 
         # gradient norm tracking
-        if not isinstance(track_grad_norm, (int, float)) and track_grad_norm != "inf":
+        if track_grad_norm != -1 and not (
+            (isinstance(track_grad_norm, (int, float)) or track_grad_norm == "inf") and float(track_grad_norm) > 0
+        ):
             raise MisconfigurationException(
-                f"`track_grad_norm` should be an int, a float or 'inf' (infinity norm). Got {track_grad_norm}."
+                f"`track_grad_norm` must be a positive number or 'inf' (infinity norm). Got {track_grad_norm}."
             )
 
         self.trainer._terminate_on_nan = terminate_on_nan

pytorch_lightning/utilities/grads.py

@@ -35,6 +35,9 @@ def grad_norm(module: Module, norm_type: Union[float, int, str]) -> Dict[str, fl
             as a single vector.
     """
     norm_type = float(norm_type)
+    if norm_type <= 0:
+        raise ValueError(f"`norm_type` must be a positive number or 'inf' (infinity norm). Got {norm_type}")
+
     norms = {
         f"grad_{norm_type}_norm_{name}": p.grad.data.norm(norm_type).item()
         for name, p in module.named_parameters()

tests/trainer/test_trainer.py

@@ -911,9 +911,10 @@ def test_invalid_terminate_on_nan(tmpdir):
         Trainer(default_root_dir=tmpdir, terminate_on_nan="False")
 
 
-def test_invalid_track_grad_norm(tmpdir):
-    with pytest.raises(MisconfigurationException, match="`track_grad_norm` should be an int, a float"):
-        Trainer(default_root_dir=tmpdir, track_grad_norm="nan")
+@pytest.mark.parametrize("track_grad_norm", [0, torch.tensor(1), "nan"])
+def test_invalid_track_grad_norm(tmpdir, track_grad_norm):
+    with pytest.raises(MisconfigurationException, match="`track_grad_norm` must be a positive number or 'inf'"):
+        Trainer(default_root_dir=tmpdir, track_grad_norm=track_grad_norm)
 
 
 @mock.patch("torch.Tensor.backward")

tests/utilities/test_grads.py

@@ -0,0 +1,78 @@
+# 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.
+from unittest.mock import Mock
+
+import pytest
+import torch
+import torch.nn as nn
+
+from pytorch_lightning.utilities import grad_norm
+
+
+@pytest.mark.parametrize(
+    "norm_type,expected",
+    [
+        (
+            1,
+            {"grad_1.0_norm_param0": 1 + 2 + 3, "grad_1.0_norm_param1": 4 + 5, "grad_1.0_norm_total": 15},
+        ),
+        (
+            2,
+            {
+                "grad_2.0_norm_param0": pow(1 + 4 + 9, 0.5),
+                "grad_2.0_norm_param1": pow(16 + 25, 0.5),
+                "grad_2.0_norm_total": pow(1 + 4 + 9 + 16 + 25, 0.5),
+            },
+        ),
+        (
+            3.14,
+            {
+                "grad_3.14_norm_param0": pow(1 + 2 ** 3.14 + 3 ** 3.14, 1 / 3.14),
+                "grad_3.14_norm_param1": pow(4 ** 3.14 + 5 ** 3.14, 1 / 3.14),
+                "grad_3.14_norm_total": pow(1 + 2 ** 3.14 + 3 ** 3.14 + 4 ** 3.14 + 5 ** 3.14, 1 / 3.14),
+            },
+        ),
+        (
+            "inf",
+            {
+                "grad_inf_norm_param0": max(1, 2, 3),
+                "grad_inf_norm_param1": max(4, 5),
+                "grad_inf_norm_total": max(1, 2, 3, 4, 5),
+            },
+        ),
+    ],
+)
+def test_grad_norm(norm_type, expected):
+    """Test utility function for computing the p-norm of individual parameter groups and norm in total."""
+
+    class Model(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.param0 = nn.Parameter(torch.rand(3))
+            self.param1 = nn.Parameter(torch.rand(2, 1))
+            self.param0.grad = torch.tensor([-1.0, 2.0, -3.0])
+            self.param1.grad = torch.tensor([[-4.0], [5.0]])
+            # param without grad should not contribute to norm
+            self.param2 = nn.Parameter(torch.rand(1))
+
+    model = Model()
+    norms = grad_norm(model, norm_type)
+    expected = {k: round(v, 4) for k, v in expected.items()}
+    assert norms == expected
+
+
+@pytest.mark.parametrize("norm_type", [-1, 0])
+def test_grad_norm_invalid_norm_type(norm_type):
+    with pytest.raises(ValueError, match="`norm_type` must be a positive number or 'inf'"):
+        grad_norm(Mock(), norm_type)