fix CPU unmap

mace/core/mace.cc

@@ -33,6 +33,35 @@
 #endif  // MACE_ENABLE_HEXAGON
 
 namespace mace {
+namespace {
+const unsigned char *LoadModelData(const std::string &model_data_file,
+                                   const size_t &data_size) {
+  int fd = open(model_data_file.c_str(), O_RDONLY);
+  MACE_CHECK(fd >= 0, "Failed to open model data file ",
+             model_data_file, ", error code: ", strerror(errno));
+
+  const unsigned char *model_data = static_cast<const unsigned char *>(
+      mmap(nullptr, 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: ", strerror(errno));
+
+  int ret = close(fd);
+  MACE_CHECK(ret == 0, "Failed to close model data file ",
+             model_data_file, ", error code: ", strerror(errno));
+
+  return model_data;
+}
+
+void UnloadModelData(const unsigned char *model_data,
+                     const size_t &data_size) {
+  MACE_CHECK(model_data != nullptr && data_size > 0,
+             "model_data is null or data_size is 0");
+  int ret = munmap(const_cast<unsigned char *>(model_data),
+                   data_size);
+  MACE_CHECK(ret == 0, "Failed to unmap model data file, error code: ",
+             strerror(errno));
+}
+}  // namespace
 
 // Mace Tensor
 class MaceTensor::Impl {
@@ -97,11 +126,18 @@ class MaceEngine::Impl {
                   const std::vector<std::string> &output_nodes,
                   const unsigned char *model_data);
 
+  MaceStatus Init(const NetDef *net_def,
+                  const std::vector<std::string> &input_nodes,
+                  const std::vector<std::string> &output_nodes,
+                  const std::string &model_data_file);
+
   MaceStatus Run(const std::map<std::string, MaceTensor> &inputs,
                  std::map<std::string, MaceTensor> *outputs,
                  RunMetadata *run_metadata);
 
  private:
+  const unsigned char *model_data_;
+  size_t model_data_size_;
   std::shared_ptr<OperatorRegistry> op_registry_;
   DeviceType device_type_;
   std::unique_ptr<Workspace> ws_;
@@ -116,7 +152,9 @@ class MaceEngine::Impl {
 };
 
 MaceEngine::Impl::Impl(DeviceType device_type)
-    : op_registry_(new OperatorRegistry()),
+    : model_data_(nullptr),
+      model_data_size_(0),
+      op_registry_(new OperatorRegistry()),
       device_type_(device_type),
       ws_(new Workspace()),
       net_(nullptr)
@@ -190,8 +228,34 @@ MaceStatus MaceEngine::Impl::Init(
   return MaceStatus::MACE_SUCCESS;
 }
 
+MaceStatus MaceEngine::Impl::Init(
+    const NetDef *net_def,
+    const std::vector<std::string> &input_nodes,
+    const std::vector<std::string> &output_nodes,
+    const std::string &model_data_file) {
+  LOG(INFO) << "Loading Model Data";
+  for (auto &const_tensor : net_def->tensors()) {
+    model_data_size_ = std::max(
+        model_data_size_,
+        static_cast<size_t>(const_tensor.offset() +
+            const_tensor.data_size() *
+                GetEnumTypeSize(const_tensor.data_type())));
+  }
+  model_data_ = LoadModelData(model_data_file, model_data_size_);
+
+  Init(net_def, input_nodes, output_nodes, model_data_);
+
+  if (device_type_ == DeviceType::GPU || device_type_ == DeviceType::HEXAGON) {
+    UnloadModelData(model_data_, model_data_size_);
+  }
+  return MaceStatus::MACE_SUCCESS;
+}
+
 MaceEngine::Impl::~Impl() {
   LOG(INFO) << "Destroying MaceEngine";
+  if (device_type_ == DeviceType::CPU && model_data_ != nullptr) {
+    UnloadModelData(model_data_, model_data_size_);
+  }
 #ifdef MACE_ENABLE_HEXAGON
   if (device_type_ == HEXAGON) {
     if (VLOG_IS_ON(2)) {
@@ -289,6 +353,14 @@ MaceStatus MaceEngine::Init(const NetDef *net_def,
   return impl_->Init(net_def, input_nodes, output_nodes, model_data);
 }
 
+
+MaceStatus MaceEngine::Init(const NetDef *net_def,
+                            const std::vector<std::string> &input_nodes,
+                            const std::vector<std::string> &output_nodes,
+                            const std::string &model_data_file) {
+  return impl_->Init(net_def, input_nodes, output_nodes, model_data_file);
+}
+
 MaceStatus MaceEngine::Run(const std::map<std::string, MaceTensor> &inputs,
                            std::map<std::string, MaceTensor> *outputs,
                            RunMetadata *run_metadata) {
@@ -300,32 +372,6 @@ MaceStatus MaceEngine::Run(const std::map<std::string, MaceTensor> &inputs,
   return impl_->Run(inputs, outputs, nullptr);
 }
 
-const unsigned char *LoadModelData(const std::string &model_data_file,
-                                   const size_t &data_size) {
-  int fd = open(model_data_file.c_str(), O_RDONLY);
-  MACE_CHECK(fd >= 0, "Failed to open model data file ",
-             model_data_file, ", error code: ", strerror(errno));
-
-  const unsigned char *model_data = static_cast<const unsigned char *>(
-      mmap(nullptr, 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: ", strerror(errno));
-
-  int ret = close(fd);
-  MACE_CHECK(ret == 0, "Failed to close model data file ",
-             model_data_file, ", error code: ", strerror(errno));
-
-  return model_data;
-}
-
-void UnloadModelData(const unsigned char *model_data,
-                     const size_t &data_size) {
-  int ret = munmap(const_cast<unsigned char *>(model_data),
-                   data_size);
-  MACE_CHECK(ret == 0, "Failed to unmap model data file, error code: ",
-             strerror(errno));
-}
-
 MaceStatus CreateMaceEngineFromProto(
     const std::vector<unsigned char> &model_pb,
     const std::string &model_data_file,
@@ -342,26 +388,10 @@ MaceStatus CreateMaceEngineFromProto(
   std::shared_ptr<NetDef> net_def(new NetDef());
   net_def->ParseFromArray(&model_pb[0], model_pb.size());
 
-  index_t model_data_size = 0;
-  for (auto &const_tensor : net_def->tensors()) {
-    model_data_size = std::max(
-        model_data_size,
-        static_cast<index_t>(const_tensor.offset() +
-                             const_tensor.data_size() *
-                             GetEnumTypeSize(const_tensor.data_type())));
-  }
-
-  MaceStatus status;
-  const unsigned char *model_data = nullptr;
-  model_data = LoadModelData(model_data_file, model_data_size);
-
   engine->reset(new mace::MaceEngine(device_type));
-  status = (*engine)->Init(
-      net_def.get(), input_nodes, output_nodes, model_data);
+  MaceStatus status = (*engine)->Init(
+      net_def.get(), input_nodes, output_nodes, model_data_file);
 
-  if (device_type == DeviceType::GPU || device_type == DeviceType::HEXAGON) {
-    UnloadModelData(model_data, model_data_size);
-  }
   return status;
 }
 

mace/public/mace.h

@@ -107,6 +107,11 @@ class MaceEngine {
                   const std::vector<std::string> &output_nodes,
                   const unsigned char *model_data);
 
+  MaceStatus Init(const NetDef *net_def,
+                  const std::vector<std::string> &input_nodes,
+                  const std::vector<std::string> &output_nodes,
+                  const std::string &model_data_file);
+
   MaceStatus Run(const std::map<std::string, MaceTensor> &inputs,
                  std::map<std::string, MaceTensor> *outputs);
 

mace/python/tools/mace_engine_factory.h.jinja2

@@ -28,9 +28,7 @@ namespace mace {
 {% for tag in model_tags %}
 namespace {{tag}} {
 
-extern const unsigned char *LoadModelData(const std::string &model_data_file);
-
-extern void UnloadModelData(const unsigned char *model_data);
+extern const unsigned char *LoadModelData();
 
 extern const std::shared_ptr<NetDef> CreateNet();
 
@@ -66,16 +64,18 @@ MaceStatus CreateMaceEngineFromCode(
   MaceStatus status = MaceStatus::MACE_SUCCESS;
   switch (model_name_map[model_name]) {
 {% for i in range(model_tags |length) %}
-   case {{ i }}:
-    model_data =
-        mace::{{model_tags[i]}}::LoadModelData(model_data_file);
-    net_def = mace::{{model_tags[i]}}::CreateNet();
-    engine->reset(new mace::MaceEngine(device_type));
-    status = (*engine)->Init(net_def.get(), input_nodes, output_nodes, model_data);
-    if (device_type == DeviceType::GPU || device_type == DeviceType::HEXAGON) {
-      mace::{{model_tags[i]}}::UnloadModelData(model_data);
-    }
-    break;
+    case {{ i }}:
+      net_def = mace::{{model_tags[i]}}::CreateNet();
+      engine->reset(new mace::MaceEngine(device_type));
+{% if embed_model_data %}
+      (void)model_data_file;
+      model_data =
+          mace::{{model_tags[i]}}::LoadModelData();
+      status = (*engine)->Init(net_def.get(), input_nodes, output_nodes, model_data);
+{% else %}
+      status = (*engine)->Init(net_def.get(), input_nodes, output_nodes, model_data_file);
+{% endif %}
+      break;
 {% endfor %}
    default:
      status = MaceStatus::MACE_INVALID_ARGS;

mace/python/tools/mace_engine_factory_codegen.py

@@ -20,7 +20,8 @@ from jinja2 import Environment, FileSystemLoader
 FLAGS = None
 
 
-def gen_mace_engine_factory(model_tags, template_dir, model_type, output_dir):
+def gen_mace_engine_factory(model_tags, template_dir, model_type,
+                            embed_model_data, output_dir):
     # Create the jinja2 environment.
     j2_env = Environment(
         loader=FileSystemLoader(template_dir), trim_blocks=True)
@@ -28,6 +29,7 @@ def gen_mace_engine_factory(model_tags, template_dir, model_type, output_dir):
     template_name = 'mace_engine_factory.h.jinja2'
     source = j2_env.get_template(template_name).render(
         model_tags=model_tags,
+        embed_model_data=embed_model_data,
         model_type=model_type,
     )
     with open(output_dir + '/mace_engine_factory.h', "wb") as f:

mace/python/tools/tensor_data.jinja2

@@ -23,12 +23,6 @@
 #include "mace/utils/logging.h"
 
 namespace mace {
-
-extern const unsigned char *LoadModelData(const std::string &model_data_file,
-                                          const size_t &data_size);
-extern void UnloadModelData(const unsigned char *model_data,
-                            const size_t &data_size);
-
 namespace {{tag}} {
 
 {% if embed_model_data %}
@@ -37,22 +31,11 @@ alignas(4) const unsigned char model_data[{{ model_data_size }}] = {
 };
 {% endif %}
 
-const unsigned char *LoadModelData(const std::string &model_data_file) {
 {% if embed_model_data %}
-  MACE_UNUSED(model_data_file);
+const unsigned char *LoadModelData() {
   return model_data;
-{% else %}
-  return mace::LoadModelData(model_data_file, {{ model_data_size }});
-{% endif %}
 }
-
-void UnloadModelData(const unsigned char *model_data) {
-{% if not embed_model_data %}
-  mace::UnloadModelData(model_data, {{ model_data_size }});
-{% else %}
-  MACE_UNUSED(model_data);
 {% endif %}
-}
 
 }  // namespace {{tag}}
 }  // namespace mace

tools/converter.py

@@ -594,7 +594,8 @@ def convert_model(configs):
 
     sh_commands.gen_mace_engine_factory_source(
         configs[YAMLKeyword.models].keys(),
-        configs[YAMLKeyword.build_type])
+        configs[YAMLKeyword.build_type],
+        configs[YAMLKeyword.embed_model_data])
     if configs[YAMLKeyword.build_type] == BuildType.code:
         sh.cp("-f", glob.glob("mace/codegen/engine/*.h"),
               model_header_dir)

tools/sh_commands.py

@@ -360,8 +360,9 @@ def gen_encrypted_opencl_source(codegen_path="mace/codegen"):
 
 def gen_mace_engine_factory_source(model_tags,
                                    model_load_type,
+                                   embed_model_data,
                                    codegen_path="mace/codegen"):
-    print("* Genearte mace engine creator source")
+    print("* Generate mace engine creator source")
     codegen_tools_dir = "%s/engine" % codegen_path
     sh.rm("-rf", codegen_tools_dir)
     sh.mkdir("-p", codegen_tools_dir)
@@ -369,8 +370,9 @@ def gen_mace_engine_factory_source(model_tags,
         model_tags,
         "mace/python/tools",
         model_load_type,
+        embed_model_data,
         codegen_tools_dir)
-    print("Genearte mace engine creator source done!\n")
+    print("Generate mace engine creator source done!\n")
 
 
 def pull_binaries(abi, serialno, model_output_dirs,