[Torch] Fix PyTorch NMS conversion for negative scores

python/tvm/relay/frontend/pytorch.py

@@ -1857,16 +1857,18 @@ def nms(self, inputs, input_types):
  scores = inputs[1]
  iou_threshold = inputs[2]
 
+ num_boxes = _op.shape_of(scores)
+
+ # TVM NMS assumes score > 0
+ scores = scores - _op.min(scores) + _op.const(1.0)
  # Generate data with shape (1, num_anchors, 5)
  scores = AttrCvt(op_name="expand_dims", extras={"axis": -1, "num_newaxis": 1})([scores], {})
-
- # Prepare input data for get_valid_counts
  data = _op.concatenate([scores, boxes], -1)
  data = _op.expand_dims(data, 0, 1)
- # Leverage get_valid_counts to sort the data and clear invalid boxes
- ct, data, indices = get_relay_op("get_valid_counts")(
-  data, score_threshold=-1.0, id_index=-1, score_index=0
- )
+ # PyTorch NMS doesn't have score_threshold, so no need to run get_valid_count
+ indices = _op.transform.arange(_op.squeeze(num_boxes), dtype="int32")
+ indices = _op.expand_dims(indices, 0, 1)
+ ct = num_boxes
 
  # Perform Non-Maximum Suppression,
  # PyTorch NMS doesn't have parameter top_k and max_output_size

tests/python/frontend/pytorch/test_forward.py

@@ -1675,10 +1675,10 @@ def _gen_rand_inputs(num_boxes):
  boxes = torch.rand(num_boxes, box_len, dtype=torch.float) * 0.5
  boxes[:, 2] += boxes[:, 0]
  boxes[:, 3] += boxes[:, 1]
- scores = torch.rand(num_boxes, dtype=torch.float)
+ scores = torch.from_numpy(np.random.uniform(-1, 1, size=(num_boxes,)).astype(np.float32))
  return boxes, scores
 
- targets = ["llvm"] # dynamic nms does not work on gpu
+ targets = ["llvm", "cuda"]
 
  for num_boxes, iou_thres in [(10, 0.3), (100, 0.5), (500, 0.9)]:
  in_boxes, in_scores = _gen_rand_inputs(num_boxes)

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
@@ -70,7 +71,7 @@ def generate_jit_model(index):
  ]
 
  model_func = model_funcs[index]
- model = TraceWrapper(model_func(pretrained=True))
+ model = TraceWrapper(model_func(pretrained=True, rpn_pre_nms_top_n_test=200))
 
  model.eval()
  inp = torch.Tensor(np.random.uniform(0.0, 250.0, size=(1, 3, in_size, in_size)))
@@ -94,46 +95,40 @@ 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())
+
+ score_threshold = 0.9
+ print("Num boxes:", pt_res[0].cpu().numpy().shape[0])
+  print("Num valid boxes:", np.sum(pt_res[1].cpu().numpy() >= score_threshold))