feat: add dtype info

Jeffwhen · Dec 29, 2022 · dd76be4 · dd76be4
1 parent 16dfb67
commit dd76be4
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 0 deletions.
diff --git a/pipeline/interface.cpp b/pipeline/interface.cpp
@@ -255,6 +255,7 @@ blob_info_t *get_input_info(unsigned runner_id, unsigned *num)
         blob.name = net_info->input_names[i];
         bm_shape_t &s = net_info->stages[0].input_shapes[i];
         blob.num_dims = s.num_dims;
+        blob.dtype = (net_info->input_dtypes)[i];
         memcpy(blob.dims, s.dims, s.num_dims * sizeof(int));
         blob.scale = net_info->input_scales[i];
     }
@@ -277,6 +278,7 @@ blob_info_t *get_output_info(unsigned runner_id, unsigned *num)
         blob.name = net_info->output_names[i];
         bm_shape_t &s = net_info->stages[0].output_shapes[i];
         blob.num_dims = s.num_dims;
+        blob.dtype = (net_info->output_dtypes)[i];
         memcpy(blob.dims, s.dims, s.num_dims * sizeof(int));
         blob.scale = (net_info->output_scales)[i];
     }

diff --git a/pipeline/interface.h b/pipeline/interface.h
@@ -41,6 +41,7 @@ unsigned int runner_release_output(unsigned int output_num, const tensor_data_t
 struct blob_info_t {
     const char *name;
     int num_dims;
+    int dtype;
     int dims[BM_MAX_DIMS_NUM];
     float scale;
 };

diff --git a/python/tpu_perf/infer.py b/python/tpu_perf/infer.py
@@ -54,6 +54,7 @@ class BlobInfo(ct.Structure):
     _fields_ = [
         ("name", ct.c_char_p),
         ("dims_num", ct.c_int),
+        ("dtype", ct.c_int),
         ("dims", ct.c_int * 8),
         ("scale", ct.c_float)]
 
@@ -93,6 +94,7 @@ def get_input_info(self):
         for _, info in zip(range(num.value), infos):
             result[info.name.decode()] = dict(
                 scale=info.scale,
+                dtype=info.dtype,
                 shape=[info.dims[i] for i in range(info.dims_num)])
         self.__lib.release_input_info(self.runner_id, infos)
         return result
@@ -105,6 +107,7 @@ def get_output_info(self):
         for _, info in zip(range(num.value), infos):
             result[info.name.decode()] = dict(
                 scale=info.scale,
+                dtype=info.dtype,
                 shape=[info.dims[i] for i in range(info.dims_num)])
         self.__lib.release_input_info(self.runner_id, infos)
         return result