class MicroConverter: def __init__(self, model_conf, net_def, model_weights, model_name, offset16=False, write_magic=False): self.model_conf = model_conf data_type = model_conf.get(ModelKeys.data_type, mace_pb2.DT_FLOAT) # self.net_def.arg if model_conf.get(ModelKeys.quantize_schema) == "int8": data_type = mace_pb2.DT_INT8 self.net_def = MicroIoConverter.convert(net_def, data_type) self.model_weights = model_weights self.model_name = model_name self.offset16 = offset16 self.write_magic = write_magic self.code_gen = MicroCodeGen() self.np_data_type = data_type_to_np_dt(data_type, np.float32) self.gen_folder = 'micro/codegen/' util.mkdir_p(self.gen_folder) self.op_resolver = OpResolver(self.net_def, self.model_conf) def gen_code_from_model(self, model_name, pb_model, model_weights): net_def = pb_model output_dir = self.gen_folder + 'models/' + model_name + '/' shutil.rmtree(output_dir, ignore_errors=True) util.mkdir_p(output_dir) # comput mem size and mem block offset and update the net_def, # should count before ProtoConverter mem_computer = MemComputer(net_def, self.np_data_type) tensor_mem_size = mem_computer.compute() # gen the c++ NetDef struct net_def_converter = ProtoConverter(self.offset16, self.write_magic, NetDefExcludeFields) net_def_bytes = net_def_converter.proto_to_bytes(net_def) mace_check(net_def_bytes is not None, "proto_to_bytes failed.") self.code_gen.gen_net_def_data(model_name, net_def_bytes, output_dir + 'micro_net_def_data.h') # gen operator array (op_src_path_list, op_class_name_list) = \ self.op_resolver.get_op_desc_list_from_model() self.code_gen.gen_ops_data(model_name, op_src_path_list, op_class_name_list, output_dir + 'micro_ops_list.h') # gen the c++ Graph struct graph = GraphBuilder(net_def, self.op_resolver).build() graph_converter = ProtoConverter(self.offset16, self.write_magic) graph_bytes = graph_converter.proto_to_bytes(graph) self.code_gen.gen_graph_data(model_name, graph_bytes, output_dir + 'micro_graph_data.h') scratch_buffer_size = ScratchComputer(net_def, self.model_conf).compute_size() # gen micro engine config engine_data = {} engine_data['tensor_mem_size'] = tensor_mem_size engine_data['input_size'] = len(net_def.input_info) engine_data['scratch_buffer_size'] = scratch_buffer_size self.code_gen.gen_engin_config(model_name, engine_data, output_dir + 'micro_engine_config.cc') # gen micro model tensor data tensor_bytes = bytearray(model_weights) self.code_gen.gen_model_data(model_name, tensor_bytes, output_dir + 'micro_model_data.h') def gen_engine_interface_code(self, model_name): output_dir = self.gen_folder + 'engines/' + model_name + '/' shutil.rmtree(output_dir, ignore_errors=True) util.mkdir_p(output_dir) self.code_gen.gen_engine_factory( model_name, output_dir + 'micro_engine_factory.h', output_dir + 'micro_engine_factory.cc') self.code_gen.gen_engine_c_interface( model_name, output_dir + 'micro_engine_c_interface.h', output_dir + 'micro_engine_c_interface.cc') def gen_code(self): MicroOpConverter(self.net_def, self.model_weights, self.np_data_type).convert_op_params() self.gen_code_from_model(self.model_name, self.net_def, self.model_weights) self.gen_engine_interface_code(self.model_name) def package(self, tar_package_path): (op_h_path_list, op_class_name_list) = \ self.op_resolver.get_op_desc_list_from_model() all_op_header_list = [op_desc.src_path for op_desc in McSupportedOps] op_h_exclude_list = [] for op_header in all_op_header_list: if op_header not in op_h_path_list: op_h_exclude_list.append(op_header) op_cc_exclude_list = \ [op_h.replace(".h", ".cc") for op_h in op_h_exclude_list] exclude_list = ["--exclude=" + op_h for op_h in op_h_exclude_list] exclude_list.extend( ["--exclude=" + op_h for op_h in op_cc_exclude_list]) tmp_dir = "/tmp/micro" tmp_workspace_file = "WORKSPACE" os.system("mkdir -p %s && touch %s/%s" % (tmp_dir, tmp_dir, tmp_workspace_file)) tar_command = "tar --exclude=micro/tools" tar_command += " --exclude=micro/test" tar_command += " --exclude=micro/build" tar_command += " --exclude=micro/cmake" tar_command += " --exclude=micro/codegen" tar_command += " --exclude=micro/dockerfiles" tar_command += " --exclude=micro/examples" tar_command += " --exclude=micro/third_party" tar_command += " --exclude=micro/pretrained_models" tar_command += " ".join(exclude_list) tar_command += " -zcf " + tar_package_path tar_command += " micro -C %s %s" % (tmp_dir, tmp_workspace_file) os.system(tar_command)
class MicroConverter: def __init__(self, model_conf, net_def, model_weights, model_name, offset16=False, write_magic=False): self.model_conf = model_conf data_type = model_conf.get(ModelKeys.data_type, mace_pb2.DT_FLOAT) # self.net_def.arg if model_conf.get(ModelKeys.quantize_schema) == "int8": data_type = mace_pb2.DT_INT8 self.net_def = MicroIoConverter.convert(net_def, data_type) self.model_weights = model_weights self.model_name = model_name self.offset16 = offset16 self.write_magic = write_magic self.code_gen = MicroCodeGen() self.np_data_type = data_type_to_np_dt(data_type, np.float32) self.model_dir = "micro/codegen/" + model_name + "/" util.mkdir_p(self.model_dir) self.op_resolver = OpResolver(self.net_def, self.model_conf) def gen_code_from_model(self, model_name, pb_model, model_weights): net_def = pb_model # comput mem size and mem block offset and update the net_def, # should count before ProtoConverter mem_computer = MemComputer(net_def, self.np_data_type) tensor_mem_size = mem_computer.compute() # gen the c++ NetDef struct net_def_converter = ProtoConverter(self.offset16, self.write_magic, NetDefExcludeFields) net_def_bytes = net_def_converter.proto_to_bytes(net_def) mace_check(net_def_bytes is not None, "proto_to_bytes failed.") self.code_gen.gen_net_def_data(model_name, net_def_bytes, self.model_dir + 'micro_net_def_data.h') # gen operator array (op_src_path_list, op_class_name_list, scratch_buffer_size) = \ self.op_resolver.get_op_desc_list_from_model() self.code_gen.gen_ops_data(model_name, op_src_path_list, op_class_name_list, self.model_dir + 'micro_ops_list.h') # gen the c++ Graph struct graph = GraphBuilder(net_def, self.op_resolver).build() graph_converter = ProtoConverter(self.offset16, self.write_magic) graph_bytes = graph_converter.proto_to_bytes(graph) self.code_gen.gen_graph_data(model_name, graph_bytes, self.model_dir + 'micro_graph_data.h') # gen micro engine config engine_data = {} engine_data['tensor_mem_size'] = tensor_mem_size engine_data['input_size'] = len(net_def.input_info) engine_data['scratch_buffer_size'] = scratch_buffer_size self.code_gen.gen_engin_config( model_name, engine_data, self.model_dir + 'micro_engine_config.cc') # gen micro model tensor data tensor_bytes = bytearray(model_weights) self.code_gen.gen_model_data(model_name, tensor_bytes, self.model_dir + 'micro_model_data.h') def gen_engine_interface_code(self, model_name): self.code_gen.gen_engine_factory( model_name, self.model_dir + 'micro_engine_factory.h', self.model_dir + 'micro_engine_factory.cc') self.code_gen.gen_engine_c_interface( model_name, self.model_dir + 'micro_engine_c_interface.h', self.model_dir + 'micro_engine_c_interface.cc') def gen_cmake_file(self, model_name): self.code_gen.gen_cmake_file(model_name, self.model_dir + 'CMakeLists.txt') def gen_code(self): MicroOpConverter(self.net_def, self.model_weights, self.np_data_type).convert_op_params() self.gen_code_from_model(self.model_name, self.net_def, self.model_weights) self.gen_engine_interface_code(self.model_name) self.gen_cmake_file(self.model_name) def package(self, tar_package_path): tmp_dir = "/tmp/micro" tmp_workspace_file = "WORKSPACE" os.system("mkdir -p %s && touch %s/%s" % (tmp_dir, tmp_dir, tmp_workspace_file)) tar_command = "tar --exclude=micro/tools" tar_command += " --exclude=micro/test" tar_command += " --exclude=micro/build" tar_command += " --exclude=micro/cmake" tar_command += " --exclude=micro/dockerfiles" tar_command += " --exclude=micro/examples" tar_command += " --exclude=micro/third_party" tar_command += " --exclude=micro/pretrained_models" tar_command += " -zcf " + tar_package_path tar_command += " micro -C %s %s" % (tmp_dir, tmp_workspace_file) os.system(tar_command)