From ac0a69511265f1286f4bd4e837b3d5791b47f6e9 Mon Sep 17 00:00:00 2001
From: yejianwu <yejianwu@xiaomi.com>
Date: Mon, 21 May 2018 19:50:55 +0800
Subject: [PATCH] support float in gpu

---
 mace/python/tools/converter.py            |  5 ++++-
 mace/python/tools/source_converter_lib.py | 10 +++++-----
 tools/generate_model_code.sh              |  4 ++++
 tools/mace_tools.py                       | 16 +++++++++++++---
 4 files changed, 26 insertions(+), 9 deletions(-)

diff --git a/mace/python/tools/converter.py b/mace/python/tools/converter.py
index ee4b942e..5b24fbbe 100644
--- a/mace/python/tools/converter.py
+++ b/mace/python/tools/converter.py
@@ -55,8 +55,11 @@ def main(unused_args):
         FLAGS.data_type, FLAGS.runtime, FLAGS.winograd)
 
   if FLAGS.output_type == 'source':
+    gpu_data_type = ""
+    if FLAGS.runtime == 'gpu':
+      gpu_data_type = FLAGS.data_type
     source_converter_lib.convert_to_source(output_graph_def, model_checksum, FLAGS.template, FLAGS.obfuscate,
-      FLAGS.model_tag, FLAGS.output, FLAGS.runtime, FLAGS.embed_model_data)
+      FLAGS.model_tag, FLAGS.output, FLAGS.runtime, FLAGS.embed_model_data, gpu_data_type)
   else:
     with open(FLAGS.output, "wb") as f:
       f.write(output_graph_def.SerializeToString())
diff --git a/mace/python/tools/source_converter_lib.py b/mace/python/tools/source_converter_lib.py
index 16ad18d4..e5096609 100644
--- a/mace/python/tools/source_converter_lib.py
+++ b/mace/python/tools/source_converter_lib.py
@@ -78,11 +78,11 @@ def rename_tensor(net_def):
         op.output[i] = tensor_map[op.output[i]]
 
 class TensorInfo:
-  def __init__(self, id, t, runtime):
+  def __init__(self, id, t, runtime, gpu_data_type):
     self.id = id
     self.data_type = mace_pb2.DataType.Name(t.data_type)
     if t.data_type == mace_pb2.DT_FLOAT:
-      if runtime == 'gpu':
+      if runtime == 'gpu' and gpu_data_type == 'half':
         self.data_type = mace_pb2.DT_HALF
         self.data = bytearray(np.array(t.float_data).astype(np.float16).tobytes())
       else:
@@ -96,7 +96,7 @@ class TensorInfo:
 def stringfy(value):
   return ', '.join('"{0}"'.format(w) for w in value)
 
-def convert_to_source(net_def, mode_pb_checksum, template_dir, obfuscate, model_tag, output, runtime, embed_model_data):
+def convert_to_source(net_def, mode_pb_checksum, template_dir, obfuscate, model_tag, output, runtime, embed_model_data, gpu_data_type):
   if obfuscate:
     obfuscate_name(net_def)
   else:
@@ -115,7 +115,7 @@ def convert_to_source(net_def, mode_pb_checksum, template_dir, obfuscate, model_
   offset = 0
   counter = 0
   for t in net_def.tensors:
-    tensor_info = TensorInfo(counter, t, runtime)
+    tensor_info = TensorInfo(counter, t, runtime, gpu_data_type)
     # align
     if tensor_info.data_type != 'DT_UINT8' and offset % 4 != 0:
       padding = 4 - offset % 4
@@ -167,7 +167,7 @@ def convert_to_source(net_def, mode_pb_checksum, template_dir, obfuscate, model_
 
   # generate model source files
   template_name = 'model.jinja2'
-  tensors = [TensorInfo(i, net_def.tensors[i], runtime) for i in range(len(net_def.tensors))]
+  tensors = [TensorInfo(i, net_def.tensors[i], runtime, gpu_data_type) for i in range(len(net_def.tensors))]
   source = j2_env.get_template(template_name).render(
     tensors = tensors,
     net = net_def,
diff --git a/tools/generate_model_code.sh b/tools/generate_model_code.sh
index e3baa975..e3f21773 100644
--- a/tools/generate_model_code.sh
+++ b/tools/generate_model_code.sh
@@ -3,6 +3,10 @@
 CURRENT_DIR=`dirname $0`
 source ${CURRENT_DIR}/env.sh
 
+if [ $# -eq 1 ]; then
+  DATA_TYPE=$1
+fi
+
 bazel build //mace/python/tools:converter || exit 1
 rm -rf ${MODEL_CODEGEN_DIR}
 mkdir -p ${MODEL_CODEGEN_DIR}
diff --git a/tools/mace_tools.py b/tools/mace_tools.py
index d229573a..a016bd7b 100644
--- a/tools/mace_tools.py
+++ b/tools/mace_tools.py
@@ -132,8 +132,13 @@ def generate_random_input(target_soc, model_output_dir,
         else:
           shutil.copy(input_file_list[i], dst_input_file)
 
-def generate_model_code():
-  command = "bash tools/generate_model_code.sh"
+def generate_model_code(gpu_data_type):
+  data_type = ""
+  if gpu_data_type == "half":
+    data_type = "DT_HALF"
+  elif gpu_data_type == "float":
+    data_type = "DT_FLOAT"
+  command = "bash tools/generate_model_code.sh {}".format(data_type)
   run_command(command)
 
 
@@ -319,6 +324,11 @@ def parse_args():
       type=str,
       default="all",
       help="SoCs to build, comma seperated list (getprop ro.board.platform)")
+  parser.add_argument(
+      "--gpu_data_type",
+      type=str,
+      default="",
+      help="[half|float]")
   return parser.parse_known_args()
 
 def set_environment(configs):
@@ -412,7 +422,7 @@ def main(unused_args):
             model_config['input_nodes'], input_file_list)
 
         if FLAGS.mode == "build" or FLAGS.mode == "all":
-          generate_model_code()
+          generate_model_code(FLAGS.gpu_data_type)
           build_mace_run_prod(model_name, global_runtime, target_abi,
                               target_soc, model_output_dir, FLAGS.tuning)
 
-- 
GitLab