updated mask rcnn test to verify outputs and also run cuda target

tests/python/frontend/pytorch/test_object_detection.py

@@ -23,6 +23,7 @@
 
 import tvm
 
+import tvm.testing
 from tvm import relay
 from tvm.runtime.vm import VirtualMachine
 from tvm.contrib.download import download
@@ -94,46 +95,36 @@ def test_detection_models():
     download(img_url, img)
 
     input_shape = (1, 3, in_size, in_size)
-    target = "llvm"
+
     input_name = "input0"
     shape_list = [(input_name, input_shape)]
-    score_threshold = 0.9
 
     scripted_model = generate_jit_model(1)
     mod, params = relay.frontend.from_pytorch(scripted_model, shape_list)
 
-    with tvm.transform.PassContext(opt_level=3, disabled_pass=["FoldScaleAxis"]):
-        vm_exec = relay.vm.compile(mod, target=target, params=params)
-
-    ctx = tvm.cpu()
-    vm = VirtualMachine(vm_exec, ctx)
     data = process_image(img)
-    pt_res = scripted_model(data)
-    data = data.detach().numpy()
-    vm.set_input("main", **{input_name: data})
-    tvm_res = vm.run()
-
-    # Note: due to accumulated numerical error, we can't directly compare results
-    # with pytorch output. Some boxes might have a quite tiny difference in score
-    # and the order can become different. We just measure how many valid boxes
-    # there are for input image.
-    pt_scores = pt_res[1].detach().numpy().tolist()
-    tvm_scores = tvm_res[1].asnumpy().tolist()
-    num_pt_valid_scores = num_tvm_valid_scores = 0
-
-    for score in pt_scores:
-        if score >= score_threshold:
-            num_pt_valid_scores += 1
-        else:
-            break
-
-    for score in tvm_scores:
-        if score >= score_threshold:
-            num_tvm_valid_scores += 1
-        else:
-            break
-
-    assert num_pt_valid_scores == num_tvm_valid_scores, (
-        "Output mismatch: Under score threshold {}, Pytorch has {} valid "
-        "boxes while TVM has {}.".format(score_threshold, num_pt_valid_scores, num_tvm_valid_scores)
-    )
+    data_np = data.detach().numpy()
+
+    with torch.no_grad():
+        pt_res = scripted_model(data)
+
+    for target in ["llvm", "cuda"]:
+        with tvm.transform.PassContext(opt_level=3):
+            vm_exec = relay.vm.compile(mod, target=target, params=params)
+
+        ctx = tvm.context(target, 0)
+        vm = VirtualMachine(vm_exec, ctx)
+
+        vm.set_input("main", **{input_name: data_np})
+        tvm_res = vm.run()
+
+        # Bounding boxes
+        tvm.testing.assert_allclose(
+            pt_res[0].cpu().numpy(), tvm_res[0].asnumpy(), rtol=1e-5, atol=1e-5
+        )
+        # Scores
+        tvm.testing.assert_allclose(
+            pt_res[1].cpu().numpy(), tvm_res[1].asnumpy(), rtol=1e-5, atol=1e-5
+        )
+        # Class ids
+        np.testing.assert_equal(pt_res[2].cpu().numpy(), tvm_res[2].asnumpy())