def load_model(dataset_val, range_from_batch, args): import tools pb_loader = tools.import_from_path(os.path.dirname(__file__) + '/../pb') from . import h5_converter if not args.h5_path.endswith('.h5'): raise ValueError('{} should endwith *.h5'.format(args.h5_path)) custom_objects = None if args.h5_custom_objects is not None: custom_objects_file = args.h5_custom_objects custom_objects_module = tools.import_from_path(custom_objects_file) custom_objects = custom_objects_module.get_custom_objects() args.pb_path = h5_converter.convert(args.h5_path, custom_objects) return pb_loader.load_model(dataset_val, range_from_batch, args)
def load_model(dataset_val, range_from_batch, args): import tools pb_loader = tools.import_from_path(os.path.dirname(__file__) + '/../pb') from . import h5_converter if not args.h5_path.endswith('.h5'): raise ValueError('{} should endwith *.h5'.format(args.h5_path)) args.pb_path = h5_converter.convert(args.h5_path) return pb_loader.load_model(dataset_val, range_from_batch, args)
def darknet2pb(input_dir, output_pb_path='tf_model', input_node_name='input'): import tools data_loader = tools.import_from_path(os.path.join(input_dir, 'network.py')) data_loader.load_data() # get basic info info_dict = dict() with open(os.path.join(input_dir, 'info.txt'), 'r') as F_info: contents = F_info.readlines() for l in contents: if len(l) > 1 and l[0] != '#': k, v = l.strip().split(':') info_dict[k.strip()] = v.strip() input_w = (int)(info_dict['width']) input_h = (int)(info_dict['height']) channel = (int)(info_dict['channel']) dtype = info_dict['data type'] # set input placeholder inp = tf.placeholder(shape=[None, input_h, input_w, channel], dtype=dtype, name=input_node_name) yv2 = data_loader.network_forward(inp) with open(os.path.join(input_dir, 'info.txt'), 'a') as F_info: F_info.write('\n#I/O info\n') F_info.write('input node: {}\n'.format(input_node_name.split(':')[0])) F_info.write('output node: {}\n'.format(yv2.name.split(':')[0])) # # Tensorflow session sess1 = tf.Session() sess1.run(tf.global_variables_initializer()) from tensorflow.python.framework.graph_util import convert_variables_to_constants graph = convert_variables_to_constants(sess1, sess1.graph_def, [yv2.name.split(':')[0]]) tf.train.write_graph(graph, input_dir, '{}.pb'.format(output_pb_path), as_text=False) tf.reset_default_graph() return inp.name, yv2.op.name, input_w, input_h
def load_model(dataset_val, range_from_batch, args): import tools pb_loader = tools.import_from_path(os.path.dirname(__file__) + '/../pb') cfg_path = args.cfg_path weights_path = args.weights_path build_dir = tempfile.mkdtemp() decode_darknet(cfg_path, weights_path, build_dir) pb_name = 'output' dataset_input_name, tensor_output_name, input_w, input_h = darknet2pb(build_dir, pb_name, 'input') assert(args.image_w == input_w) assert(args.image_h == input_h) pb_path = os.path.join(build_dir, pb_name+'.pb') args.pb_path = pb_path args.tensor_output_name = tensor_output_name args.dataset_input_name = dataset_input_name return pb_loader.load_model(dataset_val, range_from_batch, args)
def main(): def str2bool(v): if v.lower() in ('yes', 'true', 't', 'y', '1'): return True elif v.lower() in ('no', 'false', 'f', 'n', '0'): return False else: raise argparse.ArgumentTypeError('Boolean value expected.') parser = argparse.ArgumentParser() parser.add_argument('--dataset_input_name', default='input:0') parser.add_argument('--eight_bit_mode', type=str2bool, nargs='?', const=True, default=False) parser.add_argument('--output_path', default='build/gencode_output') parser.add_argument('--output_bin_name', default='build/model.bin') parser.add_argument('--prefix', default='') parser.add_argument('--layer_start_idx', type=int, default=0) parser.add_argument('--model_loader', default='model_loader/pb') parser.add_argument('--tensorboard_mode', type=str2bool, nargs='?', const=True, default=False) parser.add_argument('--pb_path', default=None) parser.add_argument('--h5_path', default=None) parser.add_argument('--h5_custom_objects', default=None) parser.add_argument('--cfg_path', default=None) parser.add_argument('--weights_path', default=None) parser.add_argument('--tensor_input_name', default=None) parser.add_argument('--tensor_output_name', default=None) parser.add_argument('--tensor_input_min', type=float, default=0) parser.add_argument('--tensor_input_max', type=float, default=1) parser.add_argument('--tensor_input_minmax_auto', type=str2bool, nargs='?', const=True, default=True) parser.add_argument('--dataset_loader', default='dataset_loader/img_0_1.py') parser.add_argument('--dataset_pic_path', default='dataset/yolo') parser.add_argument('--dataset_path', default='dataset/yolo') parser.add_argument('--image_w', type=int, default=320) parser.add_argument('--image_h', type=int, default=240) args = parser.parse_args() eight_bit_mode = args.eight_bit_mode output_path = args.output_path output_bin_name = args.output_bin_name args.prefix = args.prefix if len(args.prefix) > 0 \ else os.path.basename(args.output_path).replace('.', '_').replace('-', '_') layer_start_idx = args.layer_start_idx model_loader = args.model_loader tensorboard_mode = args.tensorboard_mode # used in model loader pb_path = args.pb_path # used in model loader tensor_input_name = args.tensor_input_name # used in model loader tensor_output_name = args.tensor_output_name # used in model loader input_min = args.tensor_input_min # used in model loader input_max = args.tensor_input_max # used in model loader input_minmax_auto = args.tensor_input_minmax_auto # used in model loader dataset_loader = args.dataset_loader dataset_input_name = args.dataset_input_name dataset_pic_path = args.dataset_pic_path # used in dataset loader image_w = args.image_w # used in dataset loader image_h = args.image_h # used in dataset loader if ':' not in dataset_input_name: args.dataset_input_name = dataset_input_name + ':0' if output_path.endswith('.c'): output_path = output_path[:-2] dataset_loader_module = tools.import_from_path(dataset_loader) dataset_val = dataset_loader_module.load_dataset(args) model_loader_module = tools.import_from_path(model_loader) rfb = range_from_batch.RangeFromBatchMinMax() k210_layers = model_loader_module.load_model(dataset_val, rfb, args) c_file, h_file = k210_layer_to_c_code.gen_layer_list_code( k210_layers, args.eight_bit_mode, args.prefix, args.layer_start_idx) os.makedirs(os.path.dirname(output_path), exist_ok=True) with open(output_path + '.c', 'w') as of: of.write(c_file) print('generate {} finish'.format(output_path + '.c')) with open(output_path + '.h', 'w') as of: of.write(h_file) print('generate {} finish'.format(output_path + '.h')) try: output_bin = k210_layer_to_bin.gen_layer_bin(k210_layers, args.eight_bit_mode) os.makedirs(os.path.dirname(output_bin_name), exist_ok=True) with open(output_bin_name, 'wb') as of: of.write(output_bin) print('generate bin finish') except Exception as e: print(e)
def load_dataset(args): import tools img_0_1 = tools.import_from_path(os.path.dirname(__file__) + '/img_0_1.py') dataset = img_0_1.load_dataset(args) return dataset * 2 - 1