Run Yolov3 with TVM

evaluate.py

@@ -322,14 +322,45 @@ def import_onnx_ssd_resnet34(self, target="llvm", dtype="float32"):
         model_file = os.path.join(directory, "ssd-12.onnx")
         onnx_model = onnx.load(model_file)
         shape_dict = {"image": (1, 3, 1200, 1200)}
-        mod, params = relay.frontend.from_onnx(onnx_model, freeze_params=True)
+        mod, params = relay.frontend.from_onnx(onnx_model, shape_dict, freeze_params=True)
         test_files_dir = os.path.join(directory, "test_data_set_0")
 
         # downcast to float16
         mod = convert_to_dtype(mod["main"], dtype)
         dtype = "float32" if dtype == "float32" else "float16"
 
-        return (mod, params, shape_dict, dtype, target, SSDResnet34Validator(test_files_dir))
+        return (mod, params, shape_dict, dtype, target, ONNXTestSamplesValidator(test_files_dir, input_names=list(shape_dict.keys())))
+
+
+    def import_onnx_yolo_v3(self, target="llvm", dtype="float32"):
+        archive_url = "https://github.com/onnx/models/raw/main/vision/object_detection_segmentation/yolov3/model/yolov3-12.tar.gz"
+        filename = "yolov3-12.tar.gz"
+        from tvm.contrib import download
+        import onnx
+        import tarfile
+        download.download(archive_url, filename)
+        archive = tarfile.open(filename)
+        directory = "onnx_yolov3"
+        archive.extractall(directory)
+        archive.close()
+        directory = os.path.join(directory, "yolov3-12")
+        model_file = os.path.join(directory, "yolov3-12.onnx")
+        onnx_model = onnx.load(model_file)
+        shape_dict = {
+            "input_1": (1, 3, 416, 416),
+            "image_shape": (1, 2),
+        }
+        mod, params = relay.frontend.from_onnx(onnx_model, shape_dict, freeze_params=True)
+        test_files_dir = os.path.join(directory, "test_data_set_0")
+
+        # downcast to float16
+        mod = convert_to_dtype(mod["main"], dtype)
+        dtype = "float32" if dtype == "float32" else "float16"
+        print("=" * 10)
+        print(mod)
+        print("=" * 10)
+
+        return (mod, params, shape_dict, dtype, target, ONNXTestSamplesValidator(test_files_dir, input_names=list(shape_dict.keys())))
 
 
 def get_args():
@@ -970,23 +1001,26 @@ def __init__(self, input_shape, dtype="float32"):
         self.image = image
         self.inputs = { list(input_shape.keys())[0]: image }
 
-class SSDResnet34Validator(Validator):
-    def __init__(self, test_data_dir, dtype="float32"):
+class ONNXTestSamplesValidator(Validator):
+    def __init__(self, test_data_dir, input_names, dtype="float32"):
         import onnx
+        import glob
         from onnx import numpy_helper
 
         self.test_data_dir = test_data_dir
-        input_file = os.path.join(test_data_dir, "input_0.pb")
-
-        tensor = onnx.TensorProto()
-        with open(input_file, 'rb') as f:
-            tensor.ParseFromString(f.read())
-        self.inputs = { "image": numpy_helper.to_array(tensor) }
-
-        self.dtype = dtype
+        inputs_num = len(glob.glob(os.path.join(test_data_dir, 'input_*.pb')))
+        self.inputs = {}
+        for i in range(inputs_num):
+            input_file = os.path.join(test_data_dir, 'input_{}.pb'.format(i))
+            tensor = onnx.TensorProto()
+            with open(input_file, 'rb') as f:
+                tensor.ParseFromString(f.read())
+            inp = numpy_helper.to_array(tensor)
+            self.inputs[input_names[i]] = inp
 
     def Validate(self, m, ref_outputs=[], show=False):
         import onnx
+        import glob
         from onnx import numpy_helper
         # output
         if isinstance(m, tvm.runtime.vm.VirtualMachine) or isinstance(m, tvm.runtime.profiler_vm.VirtualMachineProfiler):
@@ -1002,10 +1036,12 @@ def Validate(self, m, ref_outputs=[], show=False):
                 tvm_output = m.get_output(i)
                 outputs.append(tvm_output.asnumpy())
         refs = []
-        for fn in ["output_0.pb", "output_1.pb", "output_2.pb"]:
-            name = os.path.join(self.test_data_dir, fn)
+        inputs_num = len(glob.glob(os.path.join(self.test_data_dir, 'output_*.pb')))
+        self.inputs = {}
+        for i in range(inputs_num):
+            input_file = os.path.join(self.test_data_dir, 'output_{}.pb'.format(i))
             tensor = onnx.TensorProto()
-            with open(name, 'rb') as f:
+            with open(input_file, 'rb') as f:
                 tensor.ParseFromString(f.read())
             refs.append(numpy_helper.to_array(tensor))
 
@@ -1266,13 +1302,15 @@ def _benchmark_vm(
         inputs = []
         if isinstance(validator, Validator):
             inputs = validator.GetInputDictionary()
-            for key in inputs:
-                data = tvm.nd.array(inputs[key], ctx)
-                vm.set_input("main", data)
+            data = {}
+            for k, v in inputs.items():
+                data[k] = tvm.nd.array(v, ctx)
+            vm.set_input("main", **data)
         elif isinstance(input_shape, dict):
+            data = {}
             for key in input_shape:
-                data = tvm.nd.array(np.random.normal(size=input_shape[key]).astype("float32"), ctx)
-                vm.set_input("main", data)
+                data[key] = tvm.nd.array(np.random.normal(size=input_shape[key]).astype("float32"), ctx)
+            vm.set_input("main", **data)
         else:
             data = tvm.nd.array(np.random.normal(size=input_shape).astype("float32"), ctx)
             vm.set_input("main", data)

python/tvm/topi/adreno/utils.py

@@ -544,7 +544,7 @@ def bind_data_copy(stage, axis_to_vectorize=None):
         stage.bind(block, te.thread_axis("blockIdx.z"))
         stage.bind(thread, te.thread_axis("threadIdx.z"))
     else:
-        if shape[-1] == 4:
+        if len(shape) > 0 and shape[-1] == 4:
             axes = stage.op.axis
             fused = stage.fuse(*axes[:-1])
             ftc = numpy.prod(shape[:-1])

src/relay/transforms/annotate_texture_storage.cc

@@ -421,6 +421,14 @@ class StorageInfo : private transform::DeviceAwareExprVisitor {
       if (pattern <= kCommReduce) {
         if (const auto* ttype = call->checked_type().as<TensorTypeNode>()) {
           if (ttype->shape.size() == 5) {
+            auto node0 = ttype->shape[0].as<IntImmNode>();
+            auto node1 = ttype->shape[1].as<IntImmNode>();
+            auto node2 = ttype->shape[2].as<IntImmNode>();
+            auto node3 = ttype->shape[3].as<IntImmNode>();
+            auto node4 = ttype->shape[4].as<IntImmNode>();
+            if (!node0 || !node1 || !node2 || !node3 || !node4) {
+                return false;
+            }
             supports_texture_storage = true;
           }
         }

vm_samples/run_with_python.py

@@ -26,8 +26,11 @@
 
 def get_model():
     dtype = "float32"
-    input_name = "data"
     input_shape = (relay.Any(), 3, 224, 224)
+    shape_dict = {
+        "data": (1, 3, 224, 224),
+        "weight": (64, 3, 7, 7),
+    }
     filter_shape = (64, 3, 7, 7)
     A = relay.var("data", shape=input_shape, dtype=dtype)
     B = relay.var("weight", shape=filter_shape, dtype=dtype)
@@ -46,20 +49,21 @@ def get_model():
 
     mod = relay.Function([A, B], D)
     np.random.seed(0)
-    filter_data = np.zeros(filter_shape).astype(dtype)
+    #filter_data = np.zeros(filter_shape).astype(dtype)
     params1 = {
-        "weight": tvm.nd.array(filter_data),
+        #"weight": tvm.nd.array(filter_data),
     }
     module = tvm.IRModule({})
     module["main"] = mod
-    input_shape = (1, 3, 224, 224)
-    return module, params1, input_name, input_shape
+    return module, params1, shape_dict
 
 
 def get_ssd_model():
     dtype = "float32"
-    input_name = "data"
     input_shape = (1, 3, 1200, 1200)
+    shape_dict = {
+        "data": (1, 3, 1200, 1200)
+    }
     filter_shape = (64, 3, 7, 7)
     filter_shape2 = (64, 64, 3, 3)
     A = relay.var("data", shape=input_shape, dtype=dtype)
@@ -110,7 +114,7 @@ def get_ssd_model():
     }
     module = tvm.IRModule({})
     module["main"] = mod
-    return module, params1, input_name, input_shape
+    return module, params1, shape_dict
 
 
 def download_resnet_model():
@@ -127,17 +131,14 @@ def get_resnet_model():
     model_file = download_resnet_model()
     print("Import model...")
     onnx_model = onnx.load(model_file)
-    input_name = "data"
-    input_shape = (1, 3, 224, 224)
-    #shape_dict = {
-    #    input_name: [*input_shape]
-    #}
-    #model, params = relay.frontend.from_onnx(onnx_model, shape_dict, freeze_params=True)
+    shape_dict = {
+        "data": (1, 3, 224, 224)
+    }
     model, params = relay.frontend.from_onnx(onnx_model, freeze_params=True)
     print("=" * 10)
     print(model)
     print("=" * 10)
-    return model, params, input_name, input_shape
+    return model, params, shape_dict
 
 
 def download_resnet_ssd_model():
@@ -154,17 +155,39 @@ def get_resnet_ssd_model():
     model_file = download_resnet_ssd_model()
     print("Import model...")
     onnx_model = onnx.load(model_file)
-    input_name = "image"
-    input_shape = (1, 3, 1200, 1200)
-    #shape_dict = {
-    #    input_name: [*input_shape]
-    #}
-    #model, params = relay.frontend.from_onnx(onnx_model, shape_dict, freeze_params=True)
+    shape_dict = {
+        "image": (1, 3, 1200, 1200)
+    }
     model, params = relay.frontend.from_onnx(onnx_model, freeze_params=True)
     print("=" * 10)
     print(model)
     print("=" * 10)
-    return model, params, input_name, input_shape
+    return model, params, shape_dict
+
+
+def download_onnx_yolo_model():
+    print("Download model...")
+    model_file = "onnx_yolov3.onnx"
+    model_url = "https://github.com/onnx/models/raw/main/vision/object_detection_segmentation/yolov3/model/yolov3-12.onnx"
+    if not os.path.exists(model_file):
+        import urllib.request
+        urllib.request.urlretrieve(model_url, model_file)
+    return model_file
+
+
+def get_onnx_yolo_model():
+    model_file = download_onnx_yolo_model()
+    print("Import model...")
+    onnx_model = onnx.load(model_file)
+    shape_dict = {
+        "input_1": (1, 3, 416, 416),
+        "image_shape": (1, 2),
+    }
+    model, params = relay.frontend.from_onnx(onnx_model, shape_dict, freeze_params=True)
+    print("=" * 10)
+    print(model)
+    print("=" * 10)
+    return model, params, shape_dict
 
 
 def compile_model_for_vm(name, model, params, target):
@@ -208,15 +231,16 @@ def run_model_with_vm(input_dict, lib_name):
     rlib = remote.load_module(lib_name)
     dev = remote.cl()
     vm = VirtualMachine(rlib, dev, "naive")
-    data = tvm.nd.array(list(input_dict.values())[0], dev)
-    vm.set_input("main", data)
+    data = {}
+    for k, v in input_dict.items():
+        data[k] = tvm.nd.array(v, dev)
+    vm.set_input("main", **data)
     #return vm.run()
     vm.invoke_stateful("main")
     outputs = vm.get_outputs()
     return outputs
 
 
-
 def run_model_with_ge(input_dict, lib_name):
     if RUN_ON_HOST:
         remote = rpc.LocalSession()
@@ -240,15 +264,19 @@ def run_model_with_ge(input_dict, lib_name):
 if __name__ == '__main__':
     target = Target(target_c, host=target_h)
     #name = "resnet_vm_model"
-    #model, params, input_name, input_shape = get_resnet_model()
-    #name = "my_conv2d"
-    #model, params, input_name, input_shape = get_model()
-    name = "resnet_ssd_vm_model"
-    model, params, input_name, input_shape = get_resnet_ssd_model()
+    #model, params, shape_dict = get_resnet_model()
+    name = "my_conv2d"
+    model, params, shape_dict = get_model()
+    #name = "resnet_ssd_vm_model"
+    #model, params, shape_dict = get_resnet_ssd_model()
     #name = "my_ssd_vm_model"
-    #model, params, input_name, input_shape = get_ssd_model()
-    img = np.random.rand(*input_shape).astype("float32")
-    input_dict = {input_name: img}
+    #model, params, shape_dict = get_ssd_model()
+    #name = "onnx_yolo_vm_model"
+    #model, params, shape_dict = get_onnx_yolo_model()
+    input_dict = {}
+    for k, v in shape_dict.items():
+        img = np.random.rand(*v).astype("float32")
+        input_dict[k] = img
     if USE_VM:
         _, lib_name = compile_model_for_vm(name, model, params, target)
         tvm_res = run_model_with_vm(input_dict, lib_name)