common_testing.py 5.15 KB
Newer Older
facebook-github-bot's avatar
facebook-github-bot committed
1
2
3
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.

import unittest
Nikhila Ravi's avatar
Nikhila Ravi committed
4
from pathlib import Path
Roman Shapovalov's avatar
Roman Shapovalov committed
5
from typing import Callable, Optional, Union
6
7

import numpy as np
facebook-github-bot's avatar
facebook-github-bot committed
8
import torch
Nikhila Ravi's avatar
Nikhila Ravi committed
9
10
11
12
13
14
15
16
17
from PIL import Image


def load_rgb_image(filename: str, data_dir: Union[str, Path]):
    filepath = data_dir / filename
    with Image.open(filepath) as raw_image:
        image = torch.from_numpy(np.array(raw_image) / 255.0)
    image = image.to(dtype=torch.float32)
    return image[..., :3]
facebook-github-bot's avatar
facebook-github-bot committed
18
19


Roman Shapovalov's avatar
Roman Shapovalov committed
20
21
22
TensorOrArray = Union[torch.Tensor, np.ndarray]


Nikhila Ravi's avatar
Nikhila Ravi committed
23
24
25
26
27
28
29
30
31
32
33
34
def get_random_cuda_device() -> str:
    """
    Function to get a random GPU device from the
    available devices. This is useful for testing
    that custom cuda kernels can support inputs on
    any device without having to set the device explicitly.
    """
    num_devices = torch.cuda.device_count()
    rand_device_id = torch.randint(high=num_devices, size=(1,)).item()
    return "cuda:%d" % rand_device_id


facebook-github-bot's avatar
facebook-github-bot committed
35
36
37
38
39
class TestCaseMixin(unittest.TestCase):
    def assertSeparate(self, tensor1, tensor2) -> None:
        """
        Verify that tensor1 and tensor2 have their data in distinct locations.
        """
40
        self.assertNotEqual(tensor1.storage().data_ptr(), tensor2.storage().data_ptr())
facebook-github-bot's avatar
facebook-github-bot committed
41

Georgia Gkioxari's avatar
Georgia Gkioxari committed
42
43
44
45
    def assertNotSeparate(self, tensor1, tensor2) -> None:
        """
        Verify that tensor1 and tensor2 have their data in the same locations.
        """
46
        self.assertEqual(tensor1.storage().data_ptr(), tensor2.storage().data_ptr())
Georgia Gkioxari's avatar
Georgia Gkioxari committed
47

facebook-github-bot's avatar
facebook-github-bot committed
48
49
50
51
52
53
54
55
    def assertAllSeparate(self, tensor_list) -> None:
        """
        Verify that all tensors in tensor_list have their data in
        distinct locations.
        """
        ptrs = [i.storage().data_ptr() for i in tensor_list]
        self.assertCountEqual(ptrs, set(ptrs))

Roman Shapovalov's avatar
Roman Shapovalov committed
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
    def assertNormsClose(
        self,
        input: TensorOrArray,
        other: TensorOrArray,
        norm_fn: Callable[[TensorOrArray], TensorOrArray],
        *,
        rtol: float = 1e-05,
        atol: float = 1e-08,
        equal_nan: bool = False,
        msg: Optional[str] = None,
    ) -> None:
        """
        Verifies that two tensors or arrays have the same shape and are close
            given absolute and relative tolerance; raises AssertionError otherwise.
            A custom norm function is computed before comparison. If no such pre-
            processing needed, pass `torch.abs` or, equivalently, call `assertClose`.
        Args:
            input, other: two tensors or two arrays.
            norm_fn: The function evaluates
                `all(norm_fn(input - other) <= atol + rtol * norm_fn(other))`.
                norm_fn is a tensor -> tensor function; the output has:
                    * all entries non-negative,
                    * shape defined by the input shape only.
            rtol, atol, equal_nan: as for torch.allclose.
            msg: message in case the assertion is violated.
        Note:
            Optional arguments here are all keyword-only, to avoid confusion
            with msg arguments on other assert functions.
        """

        self.assertEqual(np.shape(input), np.shape(other))

        diff = norm_fn(input - other)
        other_ = norm_fn(other)

        # We want to generalise allclose(input, output), which is essentially
        #  all(diff <= atol + rtol * other)
        # but with a sophisticated handling non-finite values.
        # We work that around by calling allclose() with the following arguments:
        # allclose(diff + other_, other_). This computes what we want because
        #  all(|diff + other_ - other_| <= atol + rtol * |other_|) ==
        #    all(|norm_fn(input - other)| <= atol + rtol * |norm_fn(other)|) ==
        #    all(norm_fn(input - other) <= atol + rtol * norm_fn(other)).

        backend = torch if torch.is_tensor(input) else np
        close = backend.allclose(
            diff + other_, other_, rtol=rtol, atol=atol, equal_nan=equal_nan
        )

        self.assertTrue(close, msg)

facebook-github-bot's avatar
facebook-github-bot committed
107
108
    def assertClose(
        self,
Roman Shapovalov's avatar
Roman Shapovalov committed
109
110
        input: TensorOrArray,
        other: TensorOrArray,
facebook-github-bot's avatar
facebook-github-bot committed
111
112
113
        *,
        rtol: float = 1e-05,
        atol: float = 1e-08,
Roman Shapovalov's avatar
Roman Shapovalov committed
114
115
        equal_nan: bool = False,
        msg: Optional[str] = None,
facebook-github-bot's avatar
facebook-github-bot committed
116
117
    ) -> None:
        """
Roman Shapovalov's avatar
Roman Shapovalov committed
118
119
120
121
        Verifies that two tensors or arrays have the same shape and are close
            given absolute and relative tolerance, i.e. checks
            `all(|input - other| <= atol + rtol * |other|)`;
            raises AssertionError otherwise.
facebook-github-bot's avatar
facebook-github-bot committed
122
123
124
        Args:
            input, other: two tensors or two arrays.
            rtol, atol, equal_nan: as for torch.allclose.
Roman Shapovalov's avatar
Roman Shapovalov committed
125
            msg: message in case the assertion is violated.
facebook-github-bot's avatar
facebook-github-bot committed
126
127
128
129
130
131
132
        Note:
            Optional arguments here are all keyword-only, to avoid confusion
            with msg arguments on other assert functions.
        """

        self.assertEqual(np.shape(input), np.shape(other))

Roman Shapovalov's avatar
Roman Shapovalov committed
133
134
135
136
137
        backend = torch if torch.is_tensor(input) else np
        close = backend.allclose(
            input, other, rtol=rtol, atol=atol, equal_nan=equal_nan
        )

Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
138
139
140
141
        if not close and msg is None:
            max_diff = backend.abs(input - other).max()
            self.fail(f"Not close. max diff {max_diff}.")

Roman Shapovalov's avatar
Roman Shapovalov committed
142
        self.assertTrue(close, msg)