common_utils.py 18 KB
Newer Older
1
2
3
4
import os
import shutil
import tempfile
import contextlib
eellison's avatar
eellison committed
5
6
7
import unittest
import argparse
import sys
8
import io
eellison's avatar
eellison committed
9
import torch
10
import warnings
eellison's avatar
eellison committed
11
import __main__
12
import random
13
import inspect
14

15
from numbers import Number
Philip Meier's avatar
Philip Meier committed
16
from torch._six import string_classes
17
from collections import OrderedDict
18
from _utils_internal import get_relative_path
19

20
21
22
import numpy as np
from PIL import Image

23
24
25
IS_PY39 = sys.version_info.major == 3 and sys.version_info.minor == 9
PY39_SEGFAULT_SKIP_MSG = "Segmentation fault with Python 3.9, see https://github.com/pytorch/vision/issues/3367"
PY39_SKIP = unittest.skipIf(IS_PY39, PY39_SEGFAULT_SKIP_MSG)
26
IN_CIRCLE_CI = os.getenv("CIRCLECI", False) == 'true'
27
28
29
IN_RE_WORKER = os.environ.get("INSIDE_RE_WORKER") is not None
IN_FBCODE = os.environ.get("IN_FBCODE_TORCHVISION") == "1"
CUDA_NOT_AVAILABLE_MSG = 'CUDA device not available'
30

31
32
33
34
35
36
37
38
39
40
41

@contextlib.contextmanager
def get_tmp_dir(src=None, **kwargs):
    tmp_dir = tempfile.mkdtemp(**kwargs)
    if src is not None:
        os.rmdir(tmp_dir)
        shutil.copytree(src, tmp_dir)
    try:
        yield tmp_dir
    finally:
        shutil.rmtree(tmp_dir)
eellison's avatar
eellison committed
42
43


44
45
46
47
48
49
def set_rng_seed(seed):
    torch.manual_seed(seed)
    random.seed(seed)
    np.random.seed(seed)


50
ACCEPT = os.getenv('EXPECTTEST_ACCEPT', '0') == '1'
51
52
TEST_WITH_SLOW = os.getenv('PYTORCH_TEST_WITH_SLOW', '0') == '1'

eellison's avatar
eellison committed
53
54
55
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

class MapNestedTensorObjectImpl(object):
    def __init__(self, tensor_map_fn):
        self.tensor_map_fn = tensor_map_fn

    def __call__(self, object):
        if isinstance(object, torch.Tensor):
            return self.tensor_map_fn(object)

        elif isinstance(object, dict):
            mapped_dict = {}
            for key, value in object.items():
                mapped_dict[self(key)] = self(value)
            return mapped_dict

        elif isinstance(object, (list, tuple)):
            mapped_iter = []
            for iter in object:
                mapped_iter.append(self(iter))
            return mapped_iter if not isinstance(object, tuple) else tuple(mapped_iter)

        else:
            return object


def map_nested_tensor_object(object, tensor_map_fn):
    impl = MapNestedTensorObjectImpl(tensor_map_fn)
    return impl(object)


83
84
85
86
87
88
89
90
def is_iterable(obj):
    try:
        iter(obj)
        return True
    except TypeError:
        return False


eellison's avatar
eellison committed
91
92
93
# adapted from TestCase in torch/test/common_utils to accept non-string
# inputs and set maximum binary size
class TestCase(unittest.TestCase):
94
95
    precision = 1e-5

96
    def _get_expected_file(self, name=None):
eellison's avatar
eellison committed
97
98
99
100
101
        # NB: we take __file__ from the module that defined the test
        # class, so we place the expect directory where the test script
        # lives, NOT where test/common_utils.py lives.
        module_id = self.__class__.__module__

102
103
104
105
106
        # Determine expected file based on environment
        expected_file_base = get_relative_path(
            os.path.realpath(sys.modules[module_id].__file__),
            "expect")

107
108
109
        # Note: for legacy reasons, the reference file names all had "ModelTest.test_" in their names
        # We hardcode it here to avoid having to re-generate the reference files
        expected_file = expected_file = os.path.join(expected_file_base, 'ModelTester.test_' + name)
eellison's avatar
eellison committed
110
111
        expected_file += "_expect.pkl"

112
113
        if not ACCEPT and not os.path.exists(expected_file):
            raise RuntimeError(
114
                f"No expect file exists for {os.path.basename(expected_file)} in {expected_file}; "
115
116
117
                "to accept the current output, re-run the failing test after setting the EXPECTTEST_ACCEPT "
                "env variable. For example: EXPECTTEST_ACCEPT=1 pytest test/test_models.py -k alexnet"
            )
118
119
120

        return expected_file

121
    def assertExpected(self, output, name, prec=None):
122
123
        r"""
        Test that a python value matches the recorded contents of a file
124
        based on a "check" name. The value must be
125
126
127
        pickable with `torch.save`. This file
        is placed in the 'expect' directory in the same directory
        as the test script. You can automatically update the recorded test
128
        output using an EXPECTTEST_ACCEPT=1 env variable.
129
        """
130
        expected_file = self._get_expected_file(name)
131
132

        if ACCEPT:
133
134
            filename = {os.path.basename(expected_file)}
            print("Accepting updated output for {}:\n\n{}".format(filename, output))
eellison's avatar
eellison committed
135
136
137
            torch.save(output, expected_file)
            MAX_PICKLE_SIZE = 50 * 1000  # 50 KB
            binary_size = os.path.getsize(expected_file)
138
139
            if binary_size > MAX_PICKLE_SIZE:
                raise RuntimeError("The output for {}, is larger than 50kb".format(filename))
eellison's avatar
eellison committed
140
        else:
141
            expected = torch.load(expected_file)
142
            self.assertEqual(output, expected, prec=prec)
eellison's avatar
eellison committed
143

144
145
146
147
148
149
150
151
152
    def assertEqual(self, x, y, prec=None, message='', allow_inf=False):
        """
        This is copied from pytorch/test/common_utils.py's TestCase.assertEqual
        """
        if isinstance(prec, str) and message == '':
            message = prec
            prec = None
        if prec is None:
            prec = self.precision
eellison's avatar
eellison committed
153

154
155
156
157
158
159
160
161
162
163
164
165
166
167
        if isinstance(x, torch.Tensor) and isinstance(y, Number):
            self.assertEqual(x.item(), y, prec=prec, message=message,
                             allow_inf=allow_inf)
        elif isinstance(y, torch.Tensor) and isinstance(x, Number):
            self.assertEqual(x, y.item(), prec=prec, message=message,
                             allow_inf=allow_inf)
        elif isinstance(x, torch.Tensor) and isinstance(y, torch.Tensor):
            def assertTensorsEqual(a, b):
                super(TestCase, self).assertEqual(a.size(), b.size(), message)
                if a.numel() > 0:
                    if (a.device.type == 'cpu' and (a.dtype == torch.float16 or a.dtype == torch.bfloat16)):
                        # CPU half and bfloat16 tensors don't have the methods we need below
                        a = a.to(torch.float32)
                    b = b.to(a)
eellison's avatar
eellison committed
168

169
170
171
172
173
174
175
176
                    if (a.dtype == torch.bool) != (b.dtype == torch.bool):
                        raise TypeError("Was expecting both tensors to be bool type.")
                    else:
                        if a.dtype == torch.bool and b.dtype == torch.bool:
                            # we want to respect precision but as bool doesn't support substraction,
                            # boolean tensor has to be converted to int
                            a = a.to(torch.int)
                            b = b.to(torch.int)
eellison's avatar
eellison committed
177

178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
                        diff = a - b
                        if a.is_floating_point():
                            # check that NaNs are in the same locations
                            nan_mask = torch.isnan(a)
                            self.assertTrue(torch.equal(nan_mask, torch.isnan(b)), message)
                            diff[nan_mask] = 0
                            # inf check if allow_inf=True
                            if allow_inf:
                                inf_mask = torch.isinf(a)
                                inf_sign = inf_mask.sign()
                                self.assertTrue(torch.equal(inf_sign, torch.isinf(b).sign()), message)
                                diff[inf_mask] = 0
                        # TODO: implement abs on CharTensor (int8)
                        if diff.is_signed() and diff.dtype != torch.int8:
                            diff = diff.abs()
                        max_err = diff.max()
                        tolerance = prec + prec * abs(a.max())
                        self.assertLessEqual(max_err, tolerance, message)
            super(TestCase, self).assertEqual(x.is_sparse, y.is_sparse, message)
            super(TestCase, self).assertEqual(x.is_quantized, y.is_quantized, message)
            if x.is_sparse:
                x = self.safeCoalesce(x)
                y = self.safeCoalesce(y)
                assertTensorsEqual(x._indices(), y._indices())
                assertTensorsEqual(x._values(), y._values())
            elif x.is_quantized and y.is_quantized:
                self.assertEqual(x.qscheme(), y.qscheme(), prec=prec,
                                 message=message, allow_inf=allow_inf)
                if x.qscheme() == torch.per_tensor_affine:
                    self.assertEqual(x.q_scale(), y.q_scale(), prec=prec,
                                     message=message, allow_inf=allow_inf)
                    self.assertEqual(x.q_zero_point(), y.q_zero_point(),
                                     prec=prec, message=message,
                                     allow_inf=allow_inf)
                elif x.qscheme() == torch.per_channel_affine:
                    self.assertEqual(x.q_per_channel_scales(), y.q_per_channel_scales(), prec=prec,
                                     message=message, allow_inf=allow_inf)
                    self.assertEqual(x.q_per_channel_zero_points(), y.q_per_channel_zero_points(),
                                     prec=prec, message=message,
                                     allow_inf=allow_inf)
                    self.assertEqual(x.q_per_channel_axis(), y.q_per_channel_axis(),
                                     prec=prec, message=message)
                self.assertEqual(x.dtype, y.dtype)
                self.assertEqual(x.int_repr().to(torch.int32),
                                 y.int_repr().to(torch.int32), prec=prec,
                                 message=message, allow_inf=allow_inf)
            else:
                assertTensorsEqual(x, y)
        elif isinstance(x, string_classes) and isinstance(y, string_classes):
            super(TestCase, self).assertEqual(x, y, message)
        elif type(x) == set and type(y) == set:
            super(TestCase, self).assertEqual(x, y, message)
        elif isinstance(x, dict) and isinstance(y, dict):
            if isinstance(x, OrderedDict) and isinstance(y, OrderedDict):
                self.assertEqual(x.items(), y.items(), prec=prec,
                                 message=message, allow_inf=allow_inf)
            else:
                self.assertEqual(set(x.keys()), set(y.keys()), prec=prec,
                                 message=message, allow_inf=allow_inf)
                key_list = list(x.keys())
                self.assertEqual([x[k] for k in key_list],
                                 [y[k] for k in key_list],
                                 prec=prec, message=message,
                                 allow_inf=allow_inf)
        elif is_iterable(x) and is_iterable(y):
            super(TestCase, self).assertEqual(len(x), len(y), message)
            for x_, y_ in zip(x, y):
                self.assertEqual(x_, y_, prec=prec, message=message,
                                 allow_inf=allow_inf)
        elif isinstance(x, bool) and isinstance(y, bool):
            super(TestCase, self).assertEqual(x, y, message)
        elif isinstance(x, Number) and isinstance(y, Number):
Philip Meier's avatar
Philip Meier committed
250
            inf = float("inf")
251
252
253
254
255
256
257
258
259
            if abs(x) == inf or abs(y) == inf:
                if allow_inf:
                    super(TestCase, self).assertEqual(x, y, message)
                else:
                    self.fail("Expected finite numeric values - x={}, y={}".format(x, y))
                return
            super(TestCase, self).assertLessEqual(abs(x - y), prec, message)
        else:
            super(TestCase, self).assertEqual(x, y, message)
eellison's avatar
eellison committed
260

261
    def check_jit_scriptable(self, nn_module, args, unwrapper=None, skip=False):
262
263
264
265
266
267
        """
        Check that a nn.Module's results in TorchScript match eager and that it
        can be exported
        """
        if not TEST_WITH_SLOW or skip:
            # TorchScript is not enabled, skip these tests
268
269
270
271
272
273
274
275
            msg = "The check_jit_scriptable test for {} was skipped. " \
                  "This test checks if the module's results in TorchScript " \
                  "match eager and that it can be exported. To run these " \
                  "tests make sure you set the environment variable " \
                  "PYTORCH_TEST_WITH_SLOW=1 and that the test is not " \
                  "manually skipped.".format(nn_module.__class__.__name__)
            warnings.warn(msg, RuntimeWarning)
            return None
eellison's avatar
eellison committed
276

277
278
279
280
281
282
283
284
285
286
        sm = torch.jit.script(nn_module)

        with freeze_rng_state():
            eager_out = nn_module(*args)

        with freeze_rng_state():
            script_out = sm(*args)
            if unwrapper:
                script_out = unwrapper(script_out)

287
        self.assertEqual(eager_out, script_out, prec=1e-4)
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
        self.assertExportImportModule(sm, args)

        return sm

    def getExportImportCopy(self, m):
        """
        Save and load a TorchScript model
        """
        buffer = io.BytesIO()
        torch.jit.save(m, buffer)
        buffer.seek(0)
        imported = torch.jit.load(buffer)
        return imported

    def assertExportImportModule(self, m, args):
        """
        Check that the results of a model are the same after saving and loading
        """
        m_import = self.getExportImportCopy(m)
        with freeze_rng_state():
            results = m(*args)
        with freeze_rng_state():
            results_from_imported = m_import(*args)
311
        self.assertEqual(results, results_from_imported, prec=3e-5)
312
313
314
315
316
317
318
319
320
321
322


@contextlib.contextmanager
def freeze_rng_state():
    rng_state = torch.get_rng_state()
    if torch.cuda.is_available():
        cuda_rng_state = torch.cuda.get_rng_state()
    yield
    if torch.cuda.is_available():
        torch.cuda.set_rng_state(cuda_rng_state)
    torch.set_rng_state(rng_state)
323
324
325
326
327


class TransformsTester(unittest.TestCase):

    def _create_data(self, height=3, width=3, channels=3, device="cpu"):
328
        tensor = torch.randint(0, 256, (channels, height, width), dtype=torch.uint8, device=device)
329
330
331
        pil_img = Image.fromarray(tensor.permute(1, 2, 0).contiguous().cpu().numpy())
        return tensor, pil_img

332
333
    def _create_data_batch(self, height=3, width=3, channels=3, num_samples=4, device="cpu"):
        batch_tensor = torch.randint(
334
            0, 256,
335
336
337
338
339
340
            (num_samples, channels, height, width),
            dtype=torch.uint8,
            device=device
        )
        return batch_tensor

341
342
343
344
345
346
347
348
349
    def compareTensorToPIL(self, tensor, pil_image, msg=None):
        np_pil_image = np.array(pil_image)
        if np_pil_image.ndim == 2:
            np_pil_image = np_pil_image[:, :, None]
        pil_tensor = torch.as_tensor(np_pil_image.transpose((2, 0, 1)))
        if msg is None:
            msg = "tensor:\n{} \ndid not equal PIL tensor:\n{}".format(tensor, pil_tensor)
        self.assertTrue(tensor.cpu().equal(pil_tensor), msg)

350
351
    def approxEqualTensorToPIL(self, tensor, pil_image, tol=1e-5, msg=None, agg_method="mean",
                               allowed_percentage_diff=None):
352
353
354
355
        np_pil_image = np.array(pil_image)
        if np_pil_image.ndim == 2:
            np_pil_image = np_pil_image[:, :, None]
        pil_tensor = torch.as_tensor(np_pil_image.transpose((2, 0, 1))).to(tensor)
356
357
358
359
360
361

        if allowed_percentage_diff is not None:
            # Assert that less than a given %age of pixels are different
            self.assertTrue(
                (tensor != pil_tensor).to(torch.float).mean() <= allowed_percentage_diff
            )
vfdev's avatar
vfdev committed
362
        # error value can be mean absolute error, max abs error
363
364
365
        # Convert to float to avoid underflow when computing absolute difference
        tensor = tensor.to(torch.float)
        pil_tensor = pil_tensor.to(torch.float)
vfdev's avatar
vfdev committed
366
        err = getattr(torch, agg_method)(torch.abs(tensor - pil_tensor)).item()
367
368
369
370
        self.assertTrue(
            err < tol,
            msg="{}: err={}, tol={}: \n{}\nvs\n{}".format(msg, err, tol, tensor[0, :10, :10], pil_tensor[0, :10, :10])
        )
371
372
373
374
375
376
377
378
379
380
381
382
383


def cycle_over(objs):
    for idx, obj in enumerate(objs):
        yield obj, objs[:idx] + objs[idx + 1:]


def int_dtypes():
    return torch.testing.integral_types()


def float_dtypes():
    return torch.testing.floating_types()
384
385
386
387
388
389
390
391


@contextlib.contextmanager
def disable_console_output():
    with contextlib.ExitStack() as stack, open(os.devnull, "w") as devnull:
        stack.enter_context(contextlib.redirect_stdout(devnull))
        stack.enter_context(contextlib.redirect_stderr(devnull))
        yield
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408


def call_args_to_kwargs_only(call_args, *callable_or_arg_names):
    callable_or_arg_name = callable_or_arg_names[0]
    if callable(callable_or_arg_name):
        argspec = inspect.getfullargspec(callable_or_arg_name)
        arg_names = argspec.args
        if isinstance(callable_or_arg_name, type):
            # remove self
            arg_names.pop(0)
    else:
        arg_names = callable_or_arg_names

    args, kwargs = call_args
    kwargs_only = kwargs.copy()
    kwargs_only.update(dict(zip(arg_names, args)))
    return kwargs_only
409
410
411


def cpu_and_gpu():
412
    # TODO: make this properly handle CircleCI
413
414
    import pytest  # noqa

415
    # ignore CPU tests in RE as they're already covered by another contbuild
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
    devices = [] if IN_RE_WORKER else ['cpu']

    if torch.cuda.is_available():
        cuda_marks = ()
    elif IN_FBCODE:
        # Dont collect cuda tests on fbcode if the machine doesnt have a GPU
        # This avoids skipping the tests. More robust would be to detect if
        # we're in sancastle instead of fbcode?
        cuda_marks = pytest.mark.dont_collect()
    else:
        cuda_marks = pytest.mark.skip(reason=CUDA_NOT_AVAILABLE_MSG)

    devices.append(pytest.param('cuda', marks=cuda_marks))

    return devices
431
432
433


def needs_cuda(test_func):
434
    # TODO: make this properly handle CircleCI
435
436
437
438
439
440
441
442
443
444
445
    import pytest  # noqa

    if IN_FBCODE and not IN_RE_WORKER:
        # We don't want to skip in fbcode, so we just don't collect
        # TODO: slightly more robust way would be to detect if we're in a sandcastle instance
        # so that the test will still be collected (and skipped) in the devvms.
        return pytest.mark.dont_collect(test_func)
    elif torch.cuda.is_available():
        return test_func
    else:
        return pytest.mark.skip(reason=CUDA_NOT_AVAILABLE_MSG)(test_func)
446
447
448
449
450
451
452
453
454
455
456


def cpu_only(test_func):
    # TODO: make this properly handle CircleCI
    import pytest  # noqa

    if IN_RE_WORKER:
        # The assumption is that all RE workers have GPUs.
        return pytest.mark.dont_collect(test_func)
    else:
        return test_func