120 lines
4.7 KiB
Python
120 lines
4.7 KiB
Python
|
import warnings
|
||
|
|
|
||
|
|
import torch
|
||
|
|
|
||
|
|
__all__ = ["detect_anomaly", "set_detect_anomaly"]
|
||
|
|
|
||
|
|
|
||
|
|
class detect_anomaly:
|
||
|
|
r"""Context-manager that enable anomaly detection for the autograd engine.
|
||
|
|
|
||
|
|
This does two things:
|
||
|
|
|
||
|
|
- Running the forward pass with detection enabled will allow the backward
|
||
|
|
pass to print the traceback of the forward operation that created the failing
|
||
|
|
backward function.
|
||
|
|
- If ``check_nan`` is ``True``, any backward computation that generate "nan"
|
||
|
|
value will raise an error. Default ``True``.
|
||
|
|
|
||
|
|
.. warning::
|
||
|
|
This mode should be enabled only for debugging as the different tests
|
||
|
|
will slow down your program execution.
|
||
|
|
|
||
|
|
Example:
|
||
|
|
|
||
|
|
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_ANOMALY)
|
||
|
|
>>> import torch
|
||
|
|
>>> from torch import autograd
|
||
|
|
>>> class MyFunc(autograd.Function):
|
||
|
|
... @staticmethod
|
||
|
|
... def forward(ctx, inp):
|
||
|
|
... return inp.clone()
|
||
|
|
... @staticmethod
|
||
|
|
... def backward(ctx, gO):
|
||
|
|
... # Error during the backward pass
|
||
|
|
... raise RuntimeError("Some error in backward")
|
||
|
|
... return gO.clone()
|
||
|
|
>>> def run_fn(a):
|
||
|
|
... out = MyFunc.apply(a)
|
||
|
|
... return out.sum()
|
||
|
|
>>> inp = torch.rand(10, 10, requires_grad=True)
|
||
|
|
>>> out = run_fn(inp)
|
||
|
|
>>> out.backward()
|
||
|
|
Traceback (most recent call last):
|
||
|
|
File "<stdin>", line 1, in <module>
|
||
|
|
File "/your/pytorch/install/torch/_tensor.py", line 93, in backward
|
||
|
|
torch.autograd.backward(self, gradient, retain_graph, create_graph)
|
||
|
|
File "/your/pytorch/install/torch/autograd/__init__.py", line 90, in backward
|
||
|
|
allow_unreachable=True) # allow_unreachable flag
|
||
|
|
File "/your/pytorch/install/torch/autograd/function.py", line 76, in apply
|
||
|
|
return self._forward_cls.backward(self, *args)
|
||
|
|
File "<stdin>", line 8, in backward
|
||
|
|
RuntimeError: Some error in backward
|
||
|
|
>>> with autograd.detect_anomaly():
|
||
|
|
... inp = torch.rand(10, 10, requires_grad=True)
|
||
|
|
... out = run_fn(inp)
|
||
|
|
... out.backward()
|
||
|
|
Traceback of forward call that caused the error:
|
||
|
|
File "tmp.py", line 53, in <module>
|
||
|
|
out = run_fn(inp)
|
||
|
|
File "tmp.py", line 44, in run_fn
|
||
|
|
out = MyFunc.apply(a)
|
||
|
|
Traceback (most recent call last):
|
||
|
|
File "<stdin>", line 4, in <module>
|
||
|
|
File "/your/pytorch/install/torch/_tensor.py", line 93, in backward
|
||
|
|
torch.autograd.backward(self, gradient, retain_graph, create_graph)
|
||
|
|
File "/your/pytorch/install/torch/autograd/__init__.py", line 90, in backward
|
||
|
|
allow_unreachable=True) # allow_unreachable flag
|
||
|
|
File "/your/pytorch/install/torch/autograd/function.py", line 76, in apply
|
||
|
|
return self._forward_cls.backward(self, *args)
|
||
|
|
File "<stdin>", line 8, in backward
|
||
|
|
RuntimeError: Some error in backward
|
||
|
|
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self, check_nan=True) -> None:
|
||
|
|
self.prev = torch.is_anomaly_enabled()
|
||
|
|
self.check_nan = check_nan
|
||
|
|
self.prev_check_nan = torch.is_anomaly_check_nan_enabled()
|
||
|
|
warnings.warn(
|
||
|
|
"Anomaly Detection has been enabled. "
|
||
|
|
"This mode will increase the runtime "
|
||
|
|
"and should only be enabled for debugging.",
|
||
|
|
stacklevel=2,
|
||
|
|
)
|
||
|
|
|
||
|
|
def __enter__(self) -> None:
|
||
|
|
torch.set_anomaly_enabled(True, self.check_nan)
|
||
|
|
|
||
|
|
def __exit__(self, *args: object) -> None:
|
||
|
|
torch.set_anomaly_enabled(self.prev, self.prev_check_nan)
|
||
|
|
|
||
|
|
|
||
|
|
class set_detect_anomaly:
|
||
|
|
r"""Context-manager that sets the anomaly detection for the autograd engine on or off.
|
||
|
|
|
||
|
|
``set_detect_anomaly`` will enable or disable the autograd anomaly detection
|
||
|
|
based on its argument :attr:`mode`.
|
||
|
|
It can be used as a context-manager or as a function.
|
||
|
|
|
||
|
|
See ``detect_anomaly`` above for details of the anomaly detection behaviour.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
mode (bool): Flag whether to enable anomaly detection (``True``),
|
||
|
|
or disable (``False``).
|
||
|
|
check_nan (bool): Flag whether to raise an error when the backward
|
||
|
|
generate "nan"
|
||
|
|
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self, mode: bool, check_nan: bool = True) -> None:
|
||
|
|
self.prev = torch.is_anomaly_enabled()
|
||
|
|
self.prev_check_nan = torch.is_anomaly_check_nan_enabled()
|
||
|
|
torch.set_anomaly_enabled(mode, check_nan)
|
||
|
|
|
||
|
|
def __enter__(self) -> None:
|
||
|
|
pass
|
||
|
|
|
||
|
|
def __exit__(self, *args: object) -> None:
|
||
|
|
torch.set_anomaly_enabled(self.prev, self.prev_check_nan)
|