[TVM] Update tvm and benchmark script (#196)

nnvm/examples/benchmark/cuda_imagenet_bench.py → nnvm/examples/benchmark/gpu_imagenet_bench.py

@@ -15,15 +15,19 @@ def tvm_callback_cuda_compile(code):
 
 def main():
     parser = argparse.ArgumentParser()
-    parser.add_argument('--model', type=str, required=True, choices=['resnet', 'mobilenet'],
-        help="The model type.")
+    parser.add_argument('--model', type=str, required=True,
+                        choices=['resnet', 'mobilenet'],
+                        help="The model type.")
+    parser.add_argument('--target', type=str, required=True,
+                        choices=['cuda', 'rocm', 'opencl', 'metal'],
+                        help="Compilation target.")
     parser.add_argument('--opt-level', type=int, default=1, help="Level of optimization.")
     parser.add_argument('--num-iter', type=int, default=1000, help="Number of iteration during benchmark.")
+    parser.add_argument('--repeat', type=int, default=1, help="Number of repeative times.")
     args = parser.parse_args()
     opt_level = args.opt_level
     num_iter = args.num_iter
-    target = "cuda"
-    ctx = tvm.gpu(0)
+    ctx = tvm.context(args.target, 0)
     batch_size = 1
     num_classes = 1000
     image_shape = (3, 224, 224)
@@ -40,11 +44,10 @@ def main():
         raise ValueError('no benchmark prepared for {}.'.format(args.model))
 
     with nnvm.compiler.build_config(opt_level=opt_level):
-        with tvm.build_config(auto_unroll_max_step=32,
-                              auto_unroll_min_depth=0,
-                              unroll_explicit=False):
+        with tvm.build_config(auto_unroll_max_step=128,
+                              unroll_explicit=(args.target != "cuda")):
             graph, lib, params = nnvm.compiler.build(
-                net, target, shape={"data": data_shape}, params=params)
+                net, args.target, shape={"data": data_shape}, params=params)
 
     data = np.random.uniform(-1, 1, size=data_shape).astype("float32")
     module = runtime.create(graph, lib, ctx)
@@ -56,11 +59,12 @@ def main():
 
     print('benchmark args: {}'.format(args))
     ftimer = module.module.time_evaluator("run", ctx, num_iter)
-    for i in range(3):
+    for i in range(args.repeat):
         prof_res = ftimer()
         print(prof_res)
-        # sleep for avoiding cpu overheat
-        time.sleep(45)
+        # sleep for avoiding device overheat
+        if i + 1 != args.repeat:
+            time.sleep(45)
 
 if __name__ == '__main__':
     main()