From fbcac14c224f81c21e813963d26a7dedfa3bccf4 Mon Sep 17 00:00:00 2001 From: Aditya Oke Date: Thu, 23 Sep 2021 22:52:18 +0530 Subject: [PATCH] Add JIT tests --- test/test_ops.py | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/test/test_ops.py b/test/test_ops.py index c64ba1fd0bb..d28126bed50 100644 --- a/test/test_ops.py +++ b/test/test_ops.py @@ -957,6 +957,14 @@ def area_check(box, expected, tolerance=1e-4): expected = torch.tensor([605113.875, 600495.1875, 592247.25]) area_check(box_tensor, expected) + def test_box_area_jit(self): + box_tensor = torch.tensor([[0, 0, 100, 100], [0, 0, 0, 0]], dtype=torch.float) + TOLERANCE = 1e-3 + expected = ops.box_area(box_tensor) + scripted_fn = torch.jit.script(ops.box_area) + scripted_area = scripted_fn(box_tensor) + torch.testing.assert_close(scripted_area, expected, rtol=0.0, atol=TOLERANCE) + class TestBoxIou: def test_iou(self): @@ -978,6 +986,14 @@ def iou_check(box, expected, tolerance=1e-4): expected = torch.tensor([[1.0, 0.9933, 0.9673], [0.9933, 1.0, 0.9737], [0.9673, 0.9737, 1.0]]) iou_check(box_tensor, expected, tolerance=0.002 if dtype == torch.float16 else 1e-4) + def test_iou_jit(self): + box_tensor = torch.tensor([[0, 0, 100, 100], [0, 0, 50, 50], [200, 200, 300, 300]], dtype=torch.float) + TOLERANCE = 1e-3 + expected = ops.box_iou(box_tensor, box_tensor) + scripted_fn = torch.jit.script(ops.box_iou) + scripted_iou = scripted_fn(box_tensor, box_tensor) + torch.testing.assert_close(scripted_iou, expected, rtol=0.0, atol=TOLERANCE) + class TestGenBoxIou: def test_gen_iou(self): @@ -999,6 +1015,14 @@ def gen_iou_check(box, expected, tolerance=1e-4): expected = torch.tensor([[1.0, 0.9933, 0.9673], [0.9933, 1.0, 0.9737], [0.9673, 0.9737, 1.0]]) gen_iou_check(box_tensor, expected, tolerance=0.002 if dtype == torch.float16 else 1e-3) + def test_giou_jit(self): + box_tensor = torch.tensor([[0, 0, 100, 100], [0, 0, 50, 50], [200, 200, 300, 300]], dtype=torch.float) + TOLERANCE = 1e-3 + expected = ops.generalized_box_iou(box_tensor, box_tensor) + scripted_fn = torch.jit.script(ops.generalized_box_iou) + scripted_iou = scripted_fn(box_tensor, box_tensor) + torch.testing.assert_close(scripted_iou, expected, rtol=0.0, atol=TOLERANCE) + class TestStochasticDepth: @pytest.mark.parametrize('p', [0.2, 0.5, 0.8])