Move NaN/Inf detection to a separate utilities file

CHANGELOG.md

@@ -9,6 +9,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added utils for NaN/Inf detection for gradients and parameters ([#6834](https://github.com/PyTorchLightning/pytorch-lightning/pull/6834/))
+
 
 - Added more explicit exception message when trying to execute `trainer.test()` or `trainer.validate()` with `fast_dev_run=True` ([#6667](https://github.com/PyTorchLightning/pytorch-lightning/pull/6667))
 
@@ -104,6 +106,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Deprecated
 
+- Deprecated `TrainerTrainingTricksMixin` in favor of a separate utilities module for NaN/Inf detection for gradients and parameters ([#6834](https://github.com/PyTorchLightning/pytorch-lightning/pull/6834/))
+
+
 - `period` has been deprecated in favor of `every_n_val_epochs` in the `ModelCheckpoint` callback ([#6146](https://github.com/PyTorchLightning/pytorch-lightning/pull/6146))
 
 

pytorch_lightning/trainer/training_loop.py

@@ -29,6 +29,7 @@
 from pytorch_lightning.utilities.distributed import rank_zero_info
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from pytorch_lightning.utilities.model_helpers import is_overridden
+from pytorch_lightning.utilities.nan import detect_nan_parameters
 from pytorch_lightning.utilities.parsing import AttributeDict
 from pytorch_lightning.utilities.warnings import WarningCache
 
@@ -636,7 +637,7 @@ def _process_closure_result(self, batch_outputs: list, opt_idx: int) -> list:
 
             # check if loss or model weights are nan
             if self.trainer.terminate_on_nan:
-                self.trainer.detect_nan_tensors(opt_closure_result.loss)
+                self._check_nan(opt_closure_result.loss)
 
             # track all the outputs across all steps
             batch_opt_idx = opt_idx if len(batch_outputs) > 1 else 0
@@ -678,7 +679,7 @@ def training_step_and_backward(self, split_batch, batch_idx, opt_idx, optimizer,
 
                     # check if loss or model weights are nan
                     if self.trainer.terminate_on_nan:
-                        self.trainer.detect_nan_tensors(result.loss)
+                        self._check_nan(result.loss)
 
                 else:
                     self.warning_cache.warn("training_step returned None if it was on purpose, ignore this warning...")
@@ -689,6 +690,12 @@ def training_step_and_backward(self, split_batch, batch_idx, opt_idx, optimizer,
 
         return result
 
+    def _check_nan(self, loss: torch.Tensor) -> None:
+        if not torch.isfinite(loss).all():
+            raise ValueError('The loss returned in `training_step` is nan or inf.')
+        model = self.trainer.lightning_module
+        detect_nan_parameters(model)
+
     def backward(self, result, optimizer, opt_idx, *args, **kwargs):
         self.trainer.dev_debugger.track_event("backward_call")
 

pytorch_lightning/trainer/training_tricks.py

@@ -19,6 +19,7 @@
 from torch import Tensor
 
 from pytorch_lightning.core.lightning import LightningModule
+from pytorch_lightning.utilities import rank_zero_deprecation
 
 EPSILON = 1e-6
 EPSILON_FP16 = 1e-5
@@ -32,12 +33,23 @@ class TrainerTrainingTricksMixin(ABC):
     lightning_module: LightningModule
 
     def print_nan_gradients(self) -> None:
+        rank_zero_deprecation(
+            "Internal: TrainerTrainingTricksMixin.print_nan_gradients is deprecated in v1.3"
+            " and will be removed in v1.5."
+            " Use `pytorch_lightning.utilities.nan.print_nan_gradients` instead."
+        )
+
         model = self.lightning_module
         for param in model.parameters():
             if (param.grad is not None) and torch.isnan(param.grad.float()).any():
                 log.info(param, param.grad)
 
     def detect_nan_tensors(self, loss: Tensor) -> None:
+        rank_zero_deprecation(
+            "Internal: TrainerTrainingTricksMixin.detect_nan_tensors is deprecated in v1.3"
+            " and will be removed in v1.5."
+            " Use `pytorch_lightning.utilities.nan.detect_nan_parameters` instead."
+        )
         model = self.lightning_module
 
         # check if loss is nan

pytorch_lightning/utilities/nan.py

@@ -0,0 +1,39 @@
+# 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.
+"""Helper functions to detect NaN/Inf values. """
+
+import logging
+
+import torch
+import torch.nn as nn
+
+log = logging.getLogger(__name__)
+
+
+def print_nan_gradients(model: nn.Module) -> None:
+    """ Iterates over model parameters and prints out parameter + gradient information if NaN. """
+    for param in model.parameters():
+        if (param.grad is not None) and torch.isnan(param.grad.float()).any():
+            log.info(param, param.grad)
+
+
+def detect_nan_parameters(model: nn.Module) -> None:
+    """ Iterates over model parameters and prints gradients if any parameter is not finite. """
+    for name, param in model.named_parameters():
+        if not torch.isfinite(param).all():
+            print_nan_gradients(model)
+            raise ValueError(
+                f'Detected nan and/or inf values in `{name}`.'
+                ' Check your forward pass for numerically unstable operations.'
+            )