diff --git a/mace/core/mace.cc b/mace/core/mace.cc index ff78d369a2d091ca2c84d7d3a134c2de6a77b82e..2150e83afcdbf05795360bfef7e55cbe4b38b08a 100644 --- a/mace/core/mace.cc +++ b/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( + 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(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 &output_nodes, const unsigned char *model_data); + MaceStatus Init(const NetDef *net_def, + const std::vector &input_nodes, + const std::vector &output_nodes, + const std::string &model_data_file); + MaceStatus Run(const std::map &inputs, std::map *outputs, RunMetadata *run_metadata); private: + const unsigned char *model_data_; + size_t model_data_size_; std::shared_ptr op_registry_; DeviceType device_type_; std::unique_ptr 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 &input_nodes, + const std::vector &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(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 &input_nodes, + const std::vector &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 &inputs, std::map *outputs, RunMetadata *run_metadata) { @@ -300,32 +372,6 @@ MaceStatus MaceEngine::Run(const std::map &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( - 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(model_data), - data_size); - MACE_CHECK(ret == 0, "Failed to unmap model data file, error code: ", - strerror(errno)); -} - MaceStatus CreateMaceEngineFromProto( const std::vector &model_pb, const std::string &model_data_file, @@ -342,26 +388,10 @@ MaceStatus CreateMaceEngineFromProto( std::shared_ptr 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(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; } diff --git a/mace/public/mace.h b/mace/public/mace.h index 059a79767d5b03fd17a85de3aeb5c96e30f8de3b..87344410c650bc217b632e305bf27407d6fb61f4 100644 --- a/mace/public/mace.h +++ b/mace/public/mace.h @@ -107,6 +107,11 @@ class MaceEngine { const std::vector &output_nodes, const unsigned char *model_data); + MaceStatus Init(const NetDef *net_def, + const std::vector &input_nodes, + const std::vector &output_nodes, + const std::string &model_data_file); + MaceStatus Run(const std::map &inputs, std::map *outputs); diff --git a/mace/python/tools/mace_engine_factory.h.jinja2 b/mace/python/tools/mace_engine_factory.h.jinja2 index 132d9e30703a33f51009cc14f9522652080a38f3..354668765d114060a8212a98775dfefc43a0fdff 100644 --- a/mace/python/tools/mace_engine_factory.h.jinja2 +++ b/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 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; diff --git a/mace/python/tools/mace_engine_factory_codegen.py b/mace/python/tools/mace_engine_factory_codegen.py index f5ecafd01197745c8c6f438d4600aa160bc7f753..36b24eaefc96f65101938d1c4358157f0aafd88a 100644 --- a/mace/python/tools/mace_engine_factory_codegen.py +++ b/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: diff --git a/mace/python/tools/tensor_data.jinja2 b/mace/python/tools/tensor_data.jinja2 index ce3ceb81d9813fb391a9d06fd846ce0ca62899a2..ca267d6d43f1dfadf4748641777c196d7f0bcab3 100644 --- a/mace/python/tools/tensor_data.jinja2 +++ b/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 diff --git a/tools/converter.py b/tools/converter.py index bfa9c80722271704e5812d48207e6b607fb1b8c4..1f0f568cb963231d9f6435f2d7833cda773567d9 100644 --- a/tools/converter.py +++ b/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) diff --git a/tools/sh_commands.py b/tools/sh_commands.py index 6499bf50a450354ada95c928e4ab3320d12ad2c3..1a48fff077b4c1ecada1e12297a8aebf54ee9d6d 100644 --- a/tools/sh_commands.py +++ b/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,