Update warnings in TrainingTricksConnector (#9595)

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

Author: 2 people

pytorch_lightning/trainer/connectors/training_trick_connector.py

@@ -24,25 +24,35 @@ def __init__(self, trainer):
 
     def on_trainer_init(
         self,
-        gradient_clip_val: float,
+        gradient_clip_val: Union[int, float],
         gradient_clip_algorithm: str,
         track_grad_norm: Union[int, float, str],
         accumulate_grad_batches: Union[int, Dict[int, int], List[list]],
         truncated_bptt_steps: Optional[int],
         terminate_on_nan: bool,
     ):
-
-        self.trainer.terminate_on_nan = terminate_on_nan
+        if not isinstance(terminate_on_nan, bool):
+            raise TypeError(f"`terminate_on_nan` should be a bool, got {terminate_on_nan}.")
 
         # gradient clipping
-        if gradient_clip_algorithm not in list(GradClipAlgorithmType):
-            raise MisconfigurationException(f"gradient_clip_algorithm should be in {list(GradClipAlgorithmType)}")
-        self.trainer.gradient_clip_val = gradient_clip_val
-        self.trainer.gradient_clip_algorithm = GradClipAlgorithmType(gradient_clip_algorithm)
+        if not isinstance(gradient_clip_val, (int, float)):
+            raise TypeError(f"`gradient_clip_val` should be an int or a float. Got {gradient_clip_val}.")
+
+        if not GradClipAlgorithmType.supported_type(gradient_clip_algorithm.lower()):
+            raise MisconfigurationException(
+                f"`gradient_clip_algorithm` {gradient_clip_algorithm} is invalid. "
+                f"Allowed algorithms: {GradClipAlgorithmType.supported_types()}."
+            )
 
         # gradient norm tracking
         if not isinstance(track_grad_norm, (int, float)) and track_grad_norm != "inf":
-            raise MisconfigurationException("track_grad_norm can be an int, a float or 'inf' (infinity norm).")
+            raise MisconfigurationException(
+                f"`track_grad_norm` should be an int, a float or 'inf' (infinity norm). Got {track_grad_norm}."
+            )
+
+        self.trainer.terminate_on_nan = terminate_on_nan
+        self.trainer.gradient_clip_val = gradient_clip_val
+        self.trainer.gradient_clip_algorithm = GradClipAlgorithmType(gradient_clip_algorithm.lower())
         self.trainer.track_grad_norm = float(track_grad_norm)
 
         # accumulated grads

pytorch_lightning/trainer/trainer.py

@@ -109,7 +109,7 @@ def __init__(
         checkpoint_callback: bool = True,
         callbacks: Optional[Union[List[Callback], Callback]] = None,
         default_root_dir: Optional[str] = None,
-        gradient_clip_val: float = 0.0,
+        gradient_clip_val: Union[int, float] = 0.0,
         gradient_clip_algorithm: str = "norm",
         process_position: int = 0,
         num_nodes: int = 1,

pytorch_lightning/utilities/enums.py

@@ -118,6 +118,14 @@ class GradClipAlgorithmType(LightningEnum):
     VALUE = "value"
     NORM = "norm"
 
+    @staticmethod
+    def supported_type(val: str) -> bool:
+        return any(x.value == val for x in GradClipAlgorithmType)
+
+    @staticmethod
+    def supported_types() -> List[str]:
+        return [x.value for x in GradClipAlgorithmType]
+
 
 class AutoRestartBatchKeys(LightningEnum):
     """

tests/trainer/test_trainer.py

@@ -800,6 +800,16 @@ def training_step(self, batch, batch_idx):
         assert torch.isfinite(param).all()
 
 
+def test_invalid_terminate_on_nan(tmpdir):
+    with pytest.raises(TypeError, match="`terminate_on_nan` should be a bool"):
+        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")
+
+
 def test_nan_params_detection(tmpdir):
     class CurrentModel(BoringModel):
         test_batch_nan = 3
@@ -1005,6 +1015,16 @@ def training_step_and_backward(split_batch, batch_idx, opt_idx, optimizer, hidde
     trainer.fit(model)
 
 
+def test_invalid_gradient_clip_value(tmpdir):
+    with pytest.raises(TypeError, match="`gradient_clip_val` should be an int or a float"):
+        Trainer(default_root_dir=tmpdir, gradient_clip_val=(1, 2))
+
+
+def test_invalid_gradient_clip_algo(tmpdir):
+    with pytest.raises(MisconfigurationException, match="`gradient_clip_algorithm` norm2 is invalid"):
+        Trainer(default_root_dir=tmpdir, gradient_clip_algorithm="norm2")
+
+
 def test_gpu_choice(tmpdir):
     trainer_options = dict(default_root_dir=tmpdir)
     # Only run if CUDA is available

tests/utilities/test_enums.py

@@ -1,4 +1,4 @@
-from pytorch_lightning.utilities import DeviceType
+from pytorch_lightning.utilities.enums import DeviceType, GradClipAlgorithmType
 
 
 def test_consistency():
@@ -9,3 +9,10 @@ def test_consistency():
     # hash cannot be case invariant
     assert DeviceType.TPU not in {"TPU", "CPU"}
     assert DeviceType.TPU in {"tpu", "CPU"}
+
+
+def test_gradient_clip_algorithms():
+    assert GradClipAlgorithmType.supported_types() == ["value", "norm"]
+    assert GradClipAlgorithmType.supported_type("norm")
+    assert GradClipAlgorithmType.supported_type("value")
+    assert not GradClipAlgorithmType.supported_type("norm2")