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)