Splice model.template to multiple files.

python/tools/model.template

@@ -10,181 +10,6 @@
 #include "mace/utils/env_time.h"
 #include "mace/utils/logging.h"
 
-{% if mode == 0 %}
-
-namespace mace {
-namespace {{tag}} {
-
-void CreateTensor{{tensor_info.id}}(std::vector<mace::ConstTensor> &tensors,
-                               const unsigned char *model_data) {
-  MACE_LATENCY_LOGGER(2, "Create tensor {{ tensor.name }}");
-  tensors.emplace_back(mace::ConstTensor(
-      {{ tensor.name|tojson }}, model_data + {{ offset }},
-      { {{ tensor.dims|join(', ') }} }, {{ tensor_info.data_type }}, {{ tensor.node_id }}));
-}
-
-}  // namespace {{tag}}
-}  // namespace mace
-
-{% elif mode == 1 %}
-{% if not embed_model_data %}
-
-#include <errno.h>
-#include <fcntl.h>
-#include <string.h>
-#include <sys/mman.h>
-#include <unistd.h>
-
-{% endif %}
-
-namespace mace {
-namespace {{tag}} {
-
-{% if embed_model_data %}
-alignas(4) const unsigned char model_data[{{ model_data_size }}] = {
-{% for d in model_data %}{{"0x%02X, " % d }}{%endfor%}
-};
-{% endif %}
-
-const unsigned char *LoadModelData(const char *model_data_file) {
-{% if embed_model_data %}
-  return model_data;
-{% else %}
-  int fd = open(model_data_file, O_RDONLY);
-  MACE_CHECK(fd >= 0, "Failed to open model data file ",
-             model_data_file, ", error code: ", errno);
-
-  const unsigned char *model_data =
-    static_cast<const unsigned char *>(mmap(nullptr, {{ model_data_size }},
-                                       PROT_READ, MAP_PRIVATE, fd, 0));
-  MACE_CHECK(model_data != MAP_FAILED, "Failed to map model data file ",
-             model_data_file, ", error code: ", errno);
-
-  int ret = close(fd);
-  MACE_CHECK(ret == 0, "Failed to close model data file ",
-             model_data_file, ", error code: ", errno);
-
-  return model_data;
-{% endif %}
-}
-
-void UnloadModelData(const unsigned char *model_data) {
-{% if not embed_model_data %}
-  int ret = munmap(const_cast<unsigned char *>(model_data),
-                   {{ model_data_size }});
-  MACE_CHECK(ret == 0, "Failed to unmap model data file, error code: ", errno);
-{% endif %}
-}
-
-}  // namespace {{tag}}
-}  // namespace mace
-
-{% elif mode == 2 %}
-
-namespace mace {
-namespace {
-
-void UpdateOp(mace::OperatorDef &op,
-              const std::string &name,
-              const std::string &type,
-              const std::vector<std::string> &inputs,
-              const std::vector<std::string> &outputs,
-              const std::vector<mace::DataType> &output_types,
-              uint32_t node_id) {
-  op.set_name(name);
-  op.set_type(type);
-  op.set_input(inputs);
-  op.set_output(outputs);
-  op.set_output_type(output_types);
-  op.set_node_id(node_id);
-}
-
-}  // namespace
-}  // namespace mace
-
-namespace mace {
-namespace {{tag}} {
-
-{% for i in range(start, end) %}
-
-void CreateOperator{{i}}(mace::OperatorDef &op) {
-  MACE_LATENCY_LOGGER(2, "Create operator {{ net.op[i].name }}");
-
-  mace::Argument *arg = nullptr;
-  {% for arg in net.op[i].arg %}
-
-  arg = op.add_arg();
-  arg->set_name({{ arg.name|tojson }});
-
-  {%- if arg.HasField('f') %}
-  arg->set_f({{ arg.f }});
-  {%- endif %}
-  {%- if arg.HasField('i') %}
-  arg->set_i({{ arg.i }});
-  {%- endif %}
-  {%- if arg.HasField('s') %}
-  arg->set_s({{ arg.s|tojson }});
-  {%- endif %}
-
-  {% if arg.floats|length != 0 %}
-  arg->set_floats({ {{ arg.floats|join(', ') }} });
-  {% endif %}
-  {% if arg.ints|length != 0 %}
-  arg->set_ints({ {{ arg.ints|join(', ') }} });
-  {% endif %}
-  {% if arg.strings|length != 0 %}
-  arg->set_strings({ {{ arg.strings|stringfy() }} });
-  {% endif %}
-  {% endfor %}
-
-  {% if net.op[i].HasField('mem_id') %}
-  op.set_mem_id({{net.op[i].mem_id}});
-  {% endif %}
-
-  {% for shape in net.op[i].output_shape %}
-	{% if shape.dims | length > 0 %}
-  op.add_output_shape(mace::OutputShape({ {{ shape.dims|join(', ') }} }));
-	{% endif %}
-  {% endfor %}
-
-  std::vector<int> output_types_int({ {{ net.op[i].output_type | join(', ') }} });
-  std::vector<mace::DataType> output_types({{ net.op[i].output_type | length }});
-  for (int k = 0; k < {{ net.op[i].output_type | length }}; ++k) {
-    output_types[k] = static_cast<mace::DataType>(output_types_int[k]);
-  }
-  UpdateOp(op, {{ net.op[i].name|tojson }}, {{ net.op[i].type|tojson}},
-          { {{ net.op[i].input|stringfy }} },
-          { {{ net.op[i].output|stringfy }} },
-          output_types,
-          {{ net.op[i].node_id }});
-
-  {% if runtime == 'dsp' %}
-    op.set_padding({{ net.op[i].padding }});
-    {% if net.op[i].node_input | length > 0 %}
-    std::vector<int> input_node_ids({ {{ net.op[i].node_input | map(attribute='node_id') | join(', ') }} });
-    std::vector<int> input_output_ports({ {{ net.op[i].node_input | map(attribute='output_port') | join(', ')}} });
-
-    for (size_t i = 0; i < {{ net.op[i].node_input | length }}; ++i) {
-      mace::NodeInput input(input_node_ids[i], input_output_ports[i]);
-      op.add_node_input(input);
-    }
-    {% endif %}
-    {% if net.op[i].out_max_byte_size | length > 0 %}
-    std::vector<int> out_max_byte_sizes {{ net.op[i].out_max_byte_size | replace('[', '{') | replace(']', '}') }};
-    for (size_t i = 0; i < {{ net.op[i].out_max_byte_size | length }}; ++i) {
-      op.add_out_max_byte_size(out_max_byte_sizes[i]);
-    }
-    {% endif %}
-  {% endif %}
-}
-
-{% endfor %}
-
-}  // namespace {{tag}}
-}  // namespace mace
-
-{% else %}
-
 namespace mace {
 namespace {{tag}} {
 
@@ -325,4 +150,3 @@ const std::string ModelChecksum() {
 
 }  // namespace {{tag}}
 }  // namespace mace
-{% endif %}

python/tools/operator.template

+//
+// Copyright (c) 2017 XiaoMi All rights reserved.
+// Generated by the mace converter.  DO NOT EDIT!
+//
+
+#include <vector>
+#include <string>
+
+#include "mace/public/mace.h"
+#include "mace/utils/env_time.h"
+#include "mace/utils/logging.h"
+
+namespace mace {
+namespace {
+
+void UpdateOp(mace::OperatorDef &op,
+              const std::string &name,
+              const std::string &type,
+              const std::vector<std::string> &inputs,
+              const std::vector<std::string> &outputs,
+              const std::vector<mace::DataType> &output_types,
+              uint32_t node_id) {
+  op.set_name(name);
+  op.set_type(type);
+  op.set_input(inputs);
+  op.set_output(outputs);
+  op.set_output_type(output_types);
+  op.set_node_id(node_id);
+}
+
+}  // namespace
+}  // namespace mace
+
+namespace mace {
+namespace {{tag}} {
+
+{% for i in range(start, end) %}
+
+void CreateOperator{{i}}(mace::OperatorDef &op) {
+  MACE_LATENCY_LOGGER(2, "Create operator {{ net.op[i].name }}");
+
+  mace::Argument *arg = nullptr;
+  {% for arg in net.op[i].arg %}
+
+  arg = op.add_arg();
+  arg->set_name({{ arg.name|tojson }});
+
+  {%- if arg.HasField('f') %}
+  arg->set_f({{ arg.f }});
+  {%- endif %}
+  {%- if arg.HasField('i') %}
+  arg->set_i({{ arg.i }});
+  {%- endif %}
+  {%- if arg.HasField('s') %}
+  arg->set_s({{ arg.s|tojson }});
+  {%- endif %}
+
+  {% if arg.floats|length != 0 %}
+  arg->set_floats({ {{ arg.floats|join(', ') }} });
+  {% endif %}
+  {% if arg.ints|length != 0 %}
+  arg->set_ints({ {{ arg.ints|join(', ') }} });
+  {% endif %}
+  {% if arg.strings|length != 0 %}
+  arg->set_strings({ {{ arg.strings|stringfy() }} });
+  {% endif %}
+  {% endfor %}
+
+  {% if net.op[i].HasField('mem_id') %}
+  op.set_mem_id({{net.op[i].mem_id}});
+  {% endif %}
+
+  {% for shape in net.op[i].output_shape %}
+	{% if shape.dims | length > 0 %}
+  op.add_output_shape(mace::OutputShape({ {{ shape.dims|join(', ') }} }));
+	{% endif %}
+  {% endfor %}
+
+  std::vector<int> output_types_int({ {{ net.op[i].output_type | join(', ') }} });
+  std::vector<mace::DataType> output_types({{ net.op[i].output_type | length }});
+  for (int k = 0; k < {{ net.op[i].output_type | length }}; ++k) {
+    output_types[k] = static_cast<mace::DataType>(output_types_int[k]);
+  }
+  UpdateOp(op, {{ net.op[i].name|tojson }}, {{ net.op[i].type|tojson}},
+          { {{ net.op[i].input|stringfy }} },
+          { {{ net.op[i].output|stringfy }} },
+          output_types,
+          {{ net.op[i].node_id }});
+
+  {% if runtime == 'dsp' %}
+    op.set_padding({{ net.op[i].padding }});
+    {% if net.op[i].node_input | length > 0 %}
+    std::vector<int> input_node_ids({ {{ net.op[i].node_input | map(attribute='node_id') | join(', ') }} });
+    std::vector<int> input_output_ports({ {{ net.op[i].node_input | map(attribute='output_port') | join(', ')}} });
+
+    for (size_t i = 0; i < {{ net.op[i].node_input | length }}; ++i) {
+      mace::NodeInput input(input_node_ids[i], input_output_ports[i]);
+      op.add_node_input(input);
+    }
+    {% endif %}
+    {% if net.op[i].out_max_byte_size | length > 0 %}
+    std::vector<int> out_max_byte_sizes {{ net.op[i].out_max_byte_size | replace('[', '{') | replace(']', '}') }};
+    for (size_t i = 0; i < {{ net.op[i].out_max_byte_size | length }}; ++i) {
+      op.add_out_max_byte_size(out_max_byte_sizes[i]);
+    }
+    {% endif %}
+  {% endif %}
+}
+
+{% endfor %}
+
+}  // namespace {{tag}}
+}  // namespace mace
+

python/tools/source_converter_lib.py

@@ -96,23 +96,21 @@ class TensorInfo:
 def stringfy(value):
   return ', '.join('"{0}"'.format(w) for w in value)
 
-def convert_to_source(net_def, mode_pb_checksum, template, 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):
   if obfuscate:
     obfuscate_name(net_def)
   else:
     rename_tensor(net_def)
 
   # Capture our current directory
-  template_dir = os.path.dirname(template)
-  template_name = os.path.basename(template)
   print template_dir
 
   # Create the jinja2 environment.
-  j2_env = Environment(loader=FileSystemLoader(template_dir),
-    trim_blocks=True)
+  j2_env = Environment(loader=FileSystemLoader(template_dir), trim_blocks=True)
   j2_env.filters['stringfy'] = stringfy
   output_dir = os.path.dirname(output) + '/'
   # generate tensor source files
+  template_name = 'tensor_source.template'
   model_data = []
   offset = 0
   counter = 0
@@ -127,7 +125,6 @@ def convert_to_source(net_def, mode_pb_checksum, template, obfuscate, model_tag,
       tensor_info = tensor_info,
       tensor = t,
       tag = model_tag,
-      mode = 0,
       runtime = runtime,
       offset = offset,
     )
@@ -138,9 +135,9 @@ def convert_to_source(net_def, mode_pb_checksum, template, obfuscate, model_tag,
     counter += 1
 
   # generate tensor data
+  template_name = 'tensor_data.template'
   source = j2_env.get_template(template_name).render(
     tag = model_tag,
-    mode = 1,
     embed_model_data = embed_model_data,
     model_data_size = offset,
     model_data = model_data
@@ -153,6 +150,7 @@ def convert_to_source(net_def, mode_pb_checksum, template, obfuscate, model_tag,
     f.close()
 
   # generate op source files
+  template_name = 'operator.template'
   counter = 0
   op_size = len(net_def.op)
   for start in range(0, op_size, 10):
@@ -161,7 +159,6 @@ def convert_to_source(net_def, mode_pb_checksum, template, obfuscate, model_tag,
       end = min(start+10, op_size),
       net = net_def,
       tag = model_tag,
-      mode = 2,
       runtime = runtime,
     )
     with open(output_dir + 'op' + str(counter) + '.cc', "wb") as f:
@@ -169,12 +166,12 @@ def convert_to_source(net_def, mode_pb_checksum, template, obfuscate, model_tag,
     counter += 1
 
   # generate model source files
+  template_name = 'model.template'
   tensors = [TensorInfo(i, net_def.tensors[i], runtime) for i in range(len(net_def.tensors))]
   source = j2_env.get_template(template_name).render(
     tensors = tensors,
     net = net_def,
     tag = model_tag,
-    mode = 3,
     runtime = runtime,
     model_pb_checksum = mode_pb_checksum
   )

python/tools/tensor_data.template

+//
+// Copyright (c) 2017 XiaoMi All rights reserved.
+// Generated by the mace converter.  DO NOT EDIT!
+//
+
+#include <vector>
+#include <string>
+
+#include "mace/public/mace.h"
+#include "mace/utils/env_time.h"
+#include "mace/utils/logging.h"
+
+{% if not embed_model_data %}
+
+#include <errno.h>
+#include <fcntl.h>
+#include <string.h>
+#include <sys/mman.h>
+#include <unistd.h>
+
+{% endif %}
+
+namespace mace {
+namespace {{tag}} {
+
+{% if embed_model_data %}
+alignas(4) const unsigned char model_data[{{ model_data_size }}] = {
+{% for d in model_data %}{{"0x%02X, " % d }}{%endfor%}
+};
+{% endif %}
+
+const unsigned char *LoadModelData(const char *model_data_file) {
+{% if embed_model_data %}
+  return model_data;
+{% else %}
+  int fd = open(model_data_file, O_RDONLY);
+  MACE_CHECK(fd >= 0, "Failed to open model data file ",
+             model_data_file, ", error code: ", errno);
+
+  const unsigned char *model_data =
+    static_cast<const unsigned char *>(mmap(nullptr, {{ model_data_size }},
+                                       PROT_READ, MAP_PRIVATE, fd, 0));
+  MACE_CHECK(model_data != MAP_FAILED, "Failed to map model data file ",
+             model_data_file, ", error code: ", errno);
+
+  int ret = close(fd);
+  MACE_CHECK(ret == 0, "Failed to close model data file ",
+             model_data_file, ", error code: ", errno);
+
+  return model_data;
+{% endif %}
+}
+
+void UnloadModelData(const unsigned char *model_data) {
+{% if not embed_model_data %}
+  int ret = munmap(const_cast<unsigned char *>(model_data),
+                   {{ model_data_size }});
+  MACE_CHECK(ret == 0, "Failed to unmap model data file, error code: ", errno);
+{% endif %}
+}
+
+}  // namespace {{tag}}
+}  // namespace mace
+

python/tools/tensor_source.template

+//
+// Copyright (c) 2017 XiaoMi All rights reserved.
+// Generated by the mace converter.  DO NOT EDIT!
+//
+
+#include <vector>
+#include <string>
+
+#include "mace/public/mace.h"
+#include "mace/utils/env_time.h"
+#include "mace/utils/logging.h"
+
+namespace mace {
+namespace {{tag}} {
+
+void CreateTensor{{tensor_info.id}}(std::vector<mace::ConstTensor> &tensors,
+                               const unsigned char *model_data) {
+  MACE_LATENCY_LOGGER(2, "Create tensor {{ tensor.name }}");
+  tensors.emplace_back(mace::ConstTensor(
+      {{ tensor.name|tojson }}, model_data + {{ offset }},
+      { {{ tensor.dims|join(', ') }} }, {{ tensor_info.data_type }}, {{ tensor.node_id }}));
+}
+
+}  // namespace {{tag}}
+}  // namespace mace
+