batched nms rewrite working

Author: masahi

nms-perf-issue/maskrcnn_test.py

@@ -1,6 +1,7 @@
 import tvm
 from tvm import relay
 from tvm.contrib.download import download
+from tvm.runtime.vm import VirtualMachine
 from tvm.relay.dataflow_pattern import *
 
 import numpy as np
@@ -23,7 +24,12 @@ def do_trace(model, inp):
 
 def dict_to_tuple(out_dict):
     if "masks" in out_dict.keys():
-        return out_dict["boxes"], out_dict["scores"], out_dict["labels"], out_dict["masks"]
+        return (
+            out_dict["boxes"],
+            out_dict["scores"],
+            out_dict["labels"],
+            out_dict["masks"],
+        )
     return out_dict["boxes"], out_dict["scores"], out_dict["labels"]
 
 
@@ -40,7 +46,8 @@ def forward(self, inp):
 def get_input():
     img_path = "test_street_small.jpg"
     img_url = (
-        "https://raw.githubusercontent.com/dmlc/web-data/" "master/gluoncv/detection/street_small.jpg"
+        "https://raw.githubusercontent.com/dmlc/web-data/"
+        "master/gluoncv/detection/street_small.jpg"
     )
     download(img_url, img_path)
 
@@ -52,12 +59,7 @@ def get_input():
     return img
 
 
-def batched_nms_pattern():
-    # exprs I want to extract
-    boxes = wildcard()
-    scores = wildcard()
-    idxs = wildcard()
-
+def batched_nms_pattern(boxes, scores, idxs, iou_threshold):
     one = is_constant()
     zero = is_constant()
 
@@ -90,24 +92,117 @@ def batched_nms_pattern():
     expand_dims = is_op("expand_dims")(concat)
 
     # %1842 = vision.get_valid_counts(%1841, -1f, meta[relay.attrs.GetValidCountsAttrs][1]);
-    return is_op("vision.get_valid_counts")(expand_dims, is_constant(), wildcard())
+    get_valid_counts_out = is_op("vision.get_valid_counts")(expand_dims, is_constant())
+    data = is_tuple_get_item(get_valid_counts_out, 1)
+    valid_counts = is_tuple_get_item(get_valid_counts_out, 0)
+    indices = is_tuple_get_item(get_valid_counts_out, 2)
+    return is_op("vision.non_max_suppression")(
+        data, valid_counts, indices, is_constant(), iou_threshold
+    )
 
 
-model_func = torchvision.models.detection.maskrcnn_resnet50_fpn
-model = TraceWrapper(model_func(pretrained=True))
-model.eval()
+def convert_batched_nms(boxes, scores, idxs, iou_thres):
+    scores = relay.expand_dims(scores, axis=-1, num_newaxis=1)
+    idxs = relay.expand_dims(idxs, axis=-1, num_newaxis=1)
+    idxs = relay.cast(idxs, "float32")
+    data = relay.concatenate([idxs, scores, boxes], -1)
+    data = relay.expand_dims(data, 0, 1)
+    ct, data, indices = relay.op.vision.get_valid_counts(
+        data, score_threshold=-1.0, id_index=0, score_index=1
+    )
+    top_k = max_out_size = -1
+    out = relay.op.vision.non_max_suppression(
+        data=data,
+        valid_count=ct,
+        indices=indices,
+        max_output_size=max_out_size,
+        iou_threshold=iou_thres,
+        force_suppress=True,
+        top_k=top_k,
+        coord_start=1,
+        score_index=1,
+        id_index=0,
+        return_indices=True,
+        invalid_to_bottom=False,
+    )
+    return out.tuple_value
 
-img = get_input()
-inp = torch.from_numpy(img)
 
-with torch.no_grad():
-    out = model(inp)
-    script_module = do_trace(model, inp)
+class NMSRewrite(DFPatternCallback):
+    def __init__(self):
+        super().__init__()
+        # exprs I want to extract
+        self.boxes = wildcard()
+        self.scores = wildcard()
+        self.idxs = wildcard()
+        self.iou_threshold = wildcard()
+
+        self.pattern = batched_nms_pattern(
+            self.boxes, self.scores, self.idxs, self.iou_threshold
+        )
+
+    def callback(self, pre, post, node_map):
+        print("matched")
+        boxes = node_map[self.boxes][0]
+        scores = node_map[self.scores][0]
+        idxs = node_map[self.idxs][0]
+        iou_thres = node_map[self.iou_threshold][0]
+        return convert_batched_nms(boxes, scores, idxs, iou_thres)
+
 
 input_name = "input0"
-shape_list = [(input_name, input_shape)]
-mod, params = relay.frontend.from_pytorch(script_module, shape_list)
+img = get_input()
+
+# model_func = torchvision.models.detection.maskrcnn_resnet50_fpn
+# model = TraceWrapper(model_func(pretrained=True))
+# model.eval()
+
+# inp = torch.from_numpy(img)
+
+# with torch.no_grad():
+#     out = model(inp)
+#     script_module = do_trace(model, inp)
+
+
+# shape_list = [(input_name, input_shape)]
+# mod, params = relay.frontend.from_pytorch(script_module, shape_list)
+
+# with open("maskrcnn_mod.json", "w") as fo:
+#     fo.write(tvm.ir.save_json(mod))
+# with open("maskrcnn.params", "wb") as fo:
+#     fo.write(relay.save_param_dict(params))
+# print(mod["main"])
+with open("maskrcnn_mod.json", "r") as fi:
+    mod = tvm.ir.load_json(fi.read())
+with open("maskrcnn.params", "rb") as fi:
+    params = relay.load_param_dict(fi.read())
+
+mod["main"] = rewrite(NMSRewrite(), mod["main"])
 # print(mod["main"])
 
-pat = batched_nms_pattern()
-print(pat.match(mod["main"].body))
+target = "llvm"
+
+with tvm.transform.PassContext(opt_level=3, disabled_pass=["FoldScaleAxis"]):
+    vm_exec = relay.vm.compile(mod, target=target, params=params)
+
+######################################################################
+# Inference with Relay VM
+# -----------------------
+ctx = tvm.cpu()
+vm = VirtualMachine(vm_exec, ctx)
+vm.set_input("main", **{input_name: img})
+tvm_res = vm.run()
+
+######################################################################
+# Get boxes with score larger than 0.9
+# ------------------------------------
+score_threshold = 0.9
+boxes = tvm_res[0].asnumpy().tolist()
+valid_boxes = []
+for i, score in enumerate(tvm_res[1].asnumpy().tolist()):
+    if score > score_threshold:
+        valid_boxes.append(boxes[i])
+    else:
+        break
+
+print("Get {} valid boxes".format(len(valid_boxes)))