# LICENSE HEADER MANAGED BY add-license-header # # Copyright 2018 Kornia 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 __future__ import annotations import math from typing import Any, Callable, Optional, Sequence, Union import pytest import torch from torch.autograd import gradcheck from torch.testing import assert_close as _assert_close Dtype = Union[torch.dtype, None] Tensor = torch.Tensor # {dtype: (rtol, atol)} _DTYPE_PRECISIONS = { torch.bfloat16: (7.8e-3, 7.8e-3), torch.float16: (1e-3, 1e-3), torch.float32: (1e-4, 1e-5), # TODO: Update to ~1.2e-7 # TODO: Update to ~2.3e-16 for fp64 torch.float64: (1e-5, 1e-5), # TODO: BaseTester used (1.3e-6, 1e-5), but it fails for general cases } def _default_tolerances(*inputs: Any) -> tuple[float, float]: rtols, atols = zip(*[_DTYPE_PRECISIONS.get(torch.as_tensor(input_).dtype, (0.0, 0.0)) for input_ in inputs]) return max(rtols), max(atols) def assert_close( actual: Tensor, expected: Tensor, *, rtol: Optional[float] = None, atol: Optional[float] = None, **kwargs: Any ) -> None: if rtol is None and atol is None: # `torch.testing.assert_close` used different default tolerances than `torch.testing.assert_allclose`. # TODO: remove this special handling as soon as https://github.com/kornia/kornia/issues/1134 is resolved # Basically, this whole wrapper function can be removed and `torch.testing.assert_close` can be used # directly. rtol, atol = _default_tolerances(actual, expected) return _assert_close( actual, expected, rtol=rtol, atol=atol, # this is the default value for torch>=1.10, but not for torch==1.9 # TODO: remove this if kornia relies on torch>=1.10 check_stride=False, equal_nan=False, **kwargs, ) def tensor_to_gradcheck_var( tensor: Tensor, dtype: Dtype = torch.float64, requires_grad: bool = True ) -> Union[Tensor, str]: """Convert the input tensor to a valid variable to check the gradient. `gradcheck` needs 64-bit floating point and requires gradient. """ if not torch.is_tensor(tensor): raise AssertionError(type(tensor)) t = tensor.type(dtype) if t.is_floating_point(): return t.requires_grad_(requires_grad) return t class BaseTester: @staticmethod def assert_close( actual: Tensor | float, expected: Tensor | float, rtol: Optional[float] = None, atol: Optional[float] = None, low_tolerance: bool = False, ) -> None: """Asserts that `actual` and `expected` are close. Args: actual: Actual input. expected: Expected input. rtol: Relative tolerance. atol: Absolute tolerance. low_tolerance: This parameter allows to reduce tolerance. Half the decimal places. Example, 1e-4 -> 1e-2 or 1e-6 -> 1e-3 """ if hasattr(actual, "data"): actual = actual.data if hasattr(expected, "data"): expected = expected.data if (isinstance(actual, Tensor) and "xla" in actual.device.type) or ( isinstance(expected, Tensor) and "xla" in expected.device.type ): rtol, atol = 1e-2, 1e-2 if (isinstance(actual, Tensor) and isinstance(expected, Tensor)) and rtol is None and atol is None: actual_rtol, actual_atol = _DTYPE_PRECISIONS.get(actual.dtype, (0.0, 0.0)) expected_rtol, expected_atol = _DTYPE_PRECISIONS.get(expected.dtype, (0.0, 0.0)) rtol, atol = max(actual_rtol, expected_rtol), max(actual_atol, expected_atol) # halve the tolerance if `low_tolerance` is true rtol = math.sqrt(rtol) if low_tolerance else rtol atol = math.sqrt(atol) if low_tolerance else atol return assert_close(actual, expected, rtol=rtol, atol=atol) @staticmethod def gradcheck( func: Callable[..., Union[torch.Tensor, Sequence[torch.Tensor]]], inputs: Union[torch.Tensor, Sequence[Any]], *, raise_exception: bool = True, fast_mode: bool = True, requires_grad: Sequence[bool] = [], dtypes: Sequence[Dtype] = [], **kwargs: Any, ) -> bool: """It will gradcheck the function using the `torch.autograd.gradcheck` method. By default this method will pass all tensor to `tensor_to_gradcheck_var` which casts the tensor to be float64 dtype, and requires grad as True. You can overwrite which tensors should have requires grad equals True, by using a Sequence of the same length of the sequence of inputs, within the requires_grad per item. You also, can overwrite with the same mechanics the dtype using the `dtypes` parameter. """ requires_grad = requires_grad if len(requires_grad) > 0 else [True] * len(inputs) dtypes = dtypes if len(dtypes) > 0 else [torch.float64] * len(inputs) # MPS does not support float64; gradcheck requires float64, so skip on MPS _all_inputs = ( [inputs] if isinstance(inputs, torch.Tensor) else list(inputs.values() if isinstance(inputs, dict) else inputs) ) if any(isinstance(t, torch.Tensor) and t.device.type == "mps" for t in _all_inputs): pytest.skip("gradcheck requires float64 which is not supported on MPS") if isinstance(inputs, torch.Tensor): inputs = tensor_to_gradcheck_var(inputs) elif isinstance(inputs, dict): inputs = { k: tensor_to_gradcheck_var(v, d, r) if isinstance(v, torch.Tensor) else v for (k, v), d, r in zip(inputs.items(), dtypes, requires_grad) } else: inputs = [ tensor_to_gradcheck_var(i, d, r) if isinstance(i, torch.Tensor) else i for i, r, d in zip(inputs, requires_grad, dtypes) ] return gradcheck(func, inputs, raise_exception=raise_exception, fast_mode=fast_mode, **kwargs)