diff --git a/test/test_functional_tensor.py b/test/test_functional_tensor.py index 42d44dfdbd9..d2bc4c8a7bc 100644 --- a/test/test_functional_tensor.py +++ b/test/test_functional_tensor.py @@ -24,17 +24,18 @@ class Tester(TransformsTester): def setUp(self): self.device = "cpu" - def _test_fn_on_batch(self, batch_tensors, fn, **fn_kwargs): + def _test_fn_on_batch(self, batch_tensors, fn, scripted_fn_atol=1e-8, **fn_kwargs): transformed_batch = fn(batch_tensors, **fn_kwargs) for i in range(len(batch_tensors)): img_tensor = batch_tensors[i, ...] transformed_img = fn(img_tensor, **fn_kwargs) self.assertTrue(transformed_img.equal(transformed_batch[i, ...])) - scripted_fn = torch.jit.script(fn) - # scriptable function test - s_transformed_batch = scripted_fn(batch_tensors, **fn_kwargs) - self.assertTrue(transformed_batch.allclose(s_transformed_batch)) + if scripted_fn_atol >= 0: + scripted_fn = torch.jit.script(fn) + # scriptable function test + s_transformed_batch = scripted_fn(batch_tensors, **fn_kwargs) + self.assertTrue(transformed_batch.allclose(s_transformed_batch, atol=scripted_fn_atol)) def test_assert_image_tensor(self): shape = (100,) @@ -166,7 +167,7 @@ def test_rgb2hsv(self): self.assertLess(max_diff, 1e-5) s_hsv_img = scripted_fn(rgb_img) - self.assertTrue(hsv_img.allclose(s_hsv_img)) + self.assertTrue(hsv_img.allclose(s_hsv_img, atol=1e-7)) batch_tensors = self._create_data_batch(120, 100, num_samples=4, device=self.device).float() self._test_fn_on_batch(batch_tensors, F_t._rgb2hsv) @@ -348,7 +349,7 @@ def _test_adjust_fn(self, fn, fn_pil, fn_t, configs, tol=2.0 + 1e-10, agg_method atol = 1.0 self.assertTrue(adjusted_tensor.allclose(scripted_result, atol=atol), msg=msg) - self._test_fn_on_batch(batch_tensors, fn, **config) + self._test_fn_on_batch(batch_tensors, fn, scripted_fn_atol=atol, **config) def test_adjust_brightness(self): self._test_adjust_fn( @@ -822,9 +823,14 @@ def test_perspective(self): if dt is not None: batch_tensors = batch_tensors.to(dtype=dt) + # Ignore the equivalence between scripted and regular function on float16 cuda. The pixels at + # the border may be entirely different due to small rounding errors. + scripted_fn_atol = -1 if (dt == torch.float16 and self.device == "cuda") else 1e-8 + for spoints, epoints in test_configs: self._test_fn_on_batch( - batch_tensors, F.perspective, startpoints=spoints, endpoints=epoints, interpolation=NEAREST + batch_tensors, F.perspective, scripted_fn_atol=scripted_fn_atol, + startpoints=spoints, endpoints=epoints, interpolation=NEAREST ) # assert changed type warning diff --git a/test/test_models.py b/test/test_models.py index 9b26839fa0b..90855fb71df 100644 --- a/test/test_models.py +++ b/test/test_models.py @@ -65,6 +65,7 @@ def get_available_video_models(): "fcn_resnet50", "fcn_resnet101", "lraspp_mobilenet_v3_large", + "maskrcnn_resnet50_fpn", )