[FIX] Save tensor size with alignment (#6487)
authorZhi <5145158+zhiics@users.noreply.github.com>
Wed, 16 Sep 2020 02:10:27 +0000 (19:10 -0700)
committerGitHub <noreply@github.com>
Wed, 16 Sep 2020 02:10:27 +0000 (19:10 -0700)
include/tvm/runtime/ndarray.h
tests/python/frontend/tensorflow/test_forward.py

index e69d802..e5f6fac 100644 (file)
@@ -409,7 +409,7 @@ inline bool SaveDLTensor(dmlc::Stream* strm, const DLTensor* tensor) {
   strm->Write(tensor->dtype);
   int ndim = tensor->ndim;
   strm->WriteArray(tensor->shape, ndim);
-  int type_bytes = tensor->dtype.bits / 8;
+  int type_bytes = (tensor->dtype.bits + 7) / 8;
   int64_t num_elems = 1;
   for (int i = 0; i < ndim; ++i) {
     num_elems *= tensor->shape[i];
index 4940b69..5d83d45 100644 (file)
@@ -115,6 +115,7 @@ def run_tvm_graph(
     layout=None,
     disabled_pass=None,
     ignore_in_shape=False,
+    serialize=False,
 ):
     """ Generic function to compile on relay and execute on tvm """
     input_data = convert_to_list(input_data)
@@ -150,6 +151,9 @@ def run_tvm_graph(
     elif mode == "vm":
         with tvm.transform.PassContext(opt_level=opt_level, disabled_pass=disabled_pass):
             vm_exec = relay.vm.compile(mod, target="llvm", params=params)
+        if serialize:
+            code, lib = vm_exec.save()
+            vm_exec = tvm.runtime.vm.Executable.load_exec(code, lib)
         vm = VirtualMachine(vm_exec, tvm.cpu())
         inputs = {}
         for e, i in zip(input_node, input_data):
@@ -2892,6 +2896,7 @@ def _test_ssd_impl():
                     out_names=out_node,
                     mode="vm",
                     disabled_pass=["FoldScaleAxis"],
+                    serialize=True,
                 )
                 for i in range(len(out_node)):
                     tvm.testing.assert_allclose(tvm_output[i], tf_output[i], rtol=1e-3, atol=1e-3)