def convert(tf_path, output_dir, version, half, align=False, input_path=None, refer_path=None, debug_mode: bool = False): checker.check_file_exist(tf_path) model_name = os.path.basename(tf_path) if output_dir is None or not os.path.isdir(output_dir): output_dir = os.path.dirname(tf_path) checker.check_file_exist(output_dir) model_name = model_name[:-len(".tflite")] if tflite2tnn(tf_path, output_dir, half) is False: logging.error("Oh No, tflite2tnn failed :(\n") sys.exit(return_code.CONVERT_FAILED) else: logging.info("Convert TensorFlowLite to TNN model succeed!\n") if version is None: version = "v1.0" if align == 'output': proto_suffix = '.tnnproto' model_suffix = '.tnnmodel' tnn_proto_name = model_name + proto_suffix tnn_model_name = model_name + model_suffix tnn_proto_path = os.path.join(output_dir, tnn_proto_name) tnn_model_path = os.path.join(output_dir, tnn_model_name) align_model.align_model(tf_path, tnn_proto_path, tnn_model_path, input_path, refer_path, None, True, debug_mode=debug_mode)
def convert(tf_path, input_names, output_names, output_dir, version, optimize, half, align=False, not_fold_const=False, input_path=None, refer_path=None): logging.info("Converter Tensorflow to TNN model\n") checker.check_file_exist(tf_path) model_name = os.path.basename(tf_path) if output_dir is None or not os.path.isdir(output_dir): output_dir = os.path.dirname(tf_path) checker.check_file_exist(output_dir) model_name = model_name[:-len(".pb")] onnx_path = os.path.join(output_dir, model_name + ".onnx") if tf2onnx(tf_path, input_names, output_names, onnx_path, not_fold_const) is False: logging.error("Oh No, tf2onnx failed :(\n") sys.exit(return_code.CONVERT_FAILED) else: logging.info("Convert TensorFlow to ONNX model succeed!\n") if version is None: version = "v1.0" checker.check_file_exist(onnx_path) onnx2tnn.convert(onnx_path, output_dir, version, optimize, half) if align is True: proto_suffix = '.tnnproto' model_suffix = '.tnnmodel' onnx_base_name = os.path.basename(onnx_path) if optimize is True: tnn_proto_name = onnx_base_name[:-len('.onnx')] + '.opt' + proto_suffix tnn_model_name = onnx_base_name[:-len('.onnx')] + '.opt' + model_suffix else: tnn_proto_name = onnx_base_name[:-len('.onnx')] + proto_suffix tnn_model_name = onnx_base_name[:-len('.onnx')] + model_suffix tnn_proto_path = os.path.join(output_dir, tnn_proto_name) tnn_model_path = os.path.join(output_dir, tnn_model_name) align_model.align_model(onnx_path, tnn_proto_path, tnn_model_path, input_path, refer_path)
def convert(proto_path, model_path, output_dir, version, optimize, half, align=False, input_path=None, refer_path=None): logging.info("Converter Caffe to ONNX Model\n") checker.check_file_exist(proto_path) checker.check_file_exist(model_path) if output_dir is None: output_dir = os.path.dirname(proto_path) checker.check_file_exist(output_dir) proto_name = os.path.basename(proto_path) proto_name = proto_name[:-len(".prototxt")] onnx_path = os.path.join(output_dir, proto_name + ".onnx") if caffe2onnx(proto_path, model_path, onnx_path) is False: logging.error("Oh No, caff2onnx failed :(\n") sys.exit(return_code.CONVERT_FAILED) else: logging.info("Congratulations! caffe2onnx succeed!\n") if version is None: version = "v1.0" is_ssd = checker.is_ssd_model(proto_path) if is_ssd: onnx2tnn.convert(onnx_path, output_dir, version, False, half, is_ssd=True) else: onnx2tnn.convert(onnx_path, output_dir, version, optimize, half) if is_ssd and ((input_path is None) or (refer_path is None)): align = False optimize = False if align is True: proto_suffix = '.tnnproto' model_suffix = '.tnnmodel' onnx_base_name = os.path.basename(onnx_path) if optimize is True: tnn_proto_name = onnx_base_name[:-len('.onnx' )] + '.opt' + proto_suffix tnn_model_name = onnx_base_name[:-len('.onnx' )] + '.opt' + model_suffix else: tnn_proto_name = onnx_base_name[:-len('.onnx')] + proto_suffix tnn_model_name = onnx_base_name[:-len('.onnx')] + model_suffix tnn_proto_path = os.path.join(output_dir, tnn_proto_name) tnn_model_path = os.path.join(output_dir, tnn_model_name) align_model.align_model(onnx_path, tnn_proto_path, tnn_model_path, input_path, refer_path)
def align_model(onnx_path: str, tnn_proto_path: str, tnn_model_path: str, input_file_path: str=None, refer_path: str = None, input_names: str = None, is_tflite: bool=False ) -> bool: """ 对 onnx 模型和 tnn 模型进行对齐. 当前支持模型: 单输入,单输出;单输入,多输出; :param onnx_path: :param tnn_proto_path: :param tnn_model_path: :return: """ logging.info("{} align model (tflite or ONNX vs TNN),please wait a moment {}\n" .format("-" * 10, "-" * 10)) checker.check_file_exist(tnn_proto_path) checker.check_file_exist(tnn_model_path) if input_names is not None: input_info = parse_input_names(input_names) # check input if input_names is not None: tnn_input_info = input_info onnx_input_info = input_info else: tnn_input_info = get_input_shape_from_tnn(tnn_proto_path) if is_tflite == True: onnx_input_info = get_input_shape_from_tflite(onnx_path) else: onnx_input_info = get_input_shape_from_onnx(onnx_path) if is_tflite == True: check_input_lite_info(onnx_input_info, tnn_input_info) else: check_input_info(onnx_input_info, tnn_input_info) if input_file_path is None: # generate data input_path = data.gene_random_data(onnx_input_info) else: if os.path.exists(input_file_path): input_path = input_file_path else: logging.error("Invalid input_file_path") sys.exit(return_code.ALIGN_FAILED) if refer_path is None: if is_tflite == True: reference_output_path = run_tflite(onnx_path, input_path, onnx_input_info) else: reference_output_path = run_onnx(onnx_path, input_path, onnx_input_info) else: if os.path.exists(refer_path): reference_output_path = refer_path else: logging.error("Invalid refer_path") sys.exit(return_code.ALIGN_FAILED) run_tnn_model_check(tnn_proto_path, tnn_model_path, input_path, reference_output_path, is_tflite) if input_file_path is None and os.path.exists(input_path): data.clean_temp_data(os.path.dirname(input_path)) if refer_path is None and os.path.exists(reference_output_path): data.clean_temp_data(reference_output_path) return True
def throw_exception(current_shape): message = "Current shape: " for item in current_shape: name, shape = item message += str(name) + ": " + str(shape) + " " logging.info( "You should use -in to specify input's name and shape. e.g.: -in name[1,3,32,32]" ) logging.info(message) exit(-1)
def check_input_info(onnx_input_info: dict, tnn_input_info: dict): if len(onnx_input_info) != len(tnn_input_info): print_not_align_message("onnx input size != tnn input size") for name, onnx_info in onnx_input_info.items(): tnn_name = convert_name.onnx_name2tnn_name(name) tnn_info = tnn_input_info[tnn_name] if check_shape_info(onnx_info, tnn_info) == True: logging.info(name + ": input shape of onnx and tnn is aligned!\n") else: logging.error("input is not align 194\n") print_not_align_message( "The {}'s shape not equal! the onnx shape:{}, tnn shape: {}\n".format(name, str(onnx_info), str(tnn_info)))
def check_input_info(onnx_input_info: dict, tnn_input_info: dict): if len(onnx_input_info) != len(tnn_input_info): print_not_align_message("onnx input size != tnn input size") for name, onnx_shape in onnx_input_info.items(): tnn_name = convert_name.onnx_name2tnn_name(name) tnn_shape = tnn_input_info[tnn_name] if type(onnx_shape[0]) is not int: onnx_shape[0] = 1 if tnn_shape != onnx_shape: print_not_align_message( "The {}'s shape not equal! the onnx shape:{}, tnn shape: {}\n". format(name, str(onnx_shape), str(tnn_shape))) logging.info("Check onnx input shape and tnn input shape align!\n")
def check_input_lite_info(onnx_input_info: dict, tnn_input_info: dict): if len(onnx_input_info) != len(tnn_input_info): print_not_align_message("tflite input size != tnn input size") for name, onnx_shape in onnx_input_info.items(): tnn_name = convert_name.onnx_name2tnn_name(name) tnn_shape = tnn_input_info[tnn_name] if type(onnx_shape[0]) is not int: onnx_shape[0] = 1 nchw = [1, 1, 1, 1] nchw[0] = onnx_shape[0] nchw[1] = onnx_shape[3] nchw[2] = onnx_shape[1] nchw[3] = onnx_shape[2] if tnn_shape != nchw: logging.info("input is not algin 216\n") print_not_align_message( "The {}'s shape not equal! the onnx shape:{}, tnn shape: {}\n". format(name, str(onnx_shape), str(tnn_shape))) logging.info("Check tflite input shape and tnn input shape align!\n")
def check_input_lite_info(onnx_input_info: dict, tnn_input_info: dict): if len(onnx_input_info) != len(tnn_input_info): print_not_align_message("tflite input size != tnn input size") for name, onnx_info in onnx_input_info.items(): tnn_name = convert_name.onnx_name2tnn_name(name) tnn_info = tnn_input_info[tnn_name] if check_shape_info(onnx_info, tnn_info): logging.info("Check tflite input shape and tnn input shape align!\n") else: logging.info("input is not align\n") print_not_align_message( "The {}'s shape not equal! the onnx shape:{}, tnn shape: {}\n".format(name, str(onnx_info), str(tnn_info))) logging.info("Check tflite input shape and tnn input shape align!\n")
def convert(onnx_path, output_dir=None, version="v1.0", optimize=True, half=False, align=False, input_path=None, refer_path=None, input_names: str = None): """ 执行 onnx 转换为 tnn 的转换指令 :parameter: onnx_path: 需要转换的 onnx 文件的路径 output_path: 生成的 tnn 文件的路径 version: 转换模型的版本号 optimize: 是否需要对模型进行优化,默认是需要进行优化 half: 是否需要转为 FP16 的模型,减小模型的大小 :return return_code :exception 执行超时 """ logging.info("Converter ONNX to TNN Model\n") checker.check_file_exist(onnx_path) ret, current_shape = checker.check_onnx_dim(onnx_path) if ret is False and current_shape is not None: if input_names is None: throw_exception(current_shape) if input_names is not None and not ("[" in input_names and "]" in input_names): throw_exception(current_shape) proto_suffix = '.tnnproto' model_suffix = '.tnnmodel' command = "python3 onnx2tnn.py " + onnx_path command = command + " -version=v1.0" checker.check_file_exist(onnx_path) if optimize is True: command = command + " -optimize=1" else: command = command + " -optimize=0" if half is True: command = command + " -half=1" else: command = command + " -half=0" if output_dir is None: output_dir = os.path.dirname(onnx_path) checker.check_file_exist(output_dir) command = command + " -o " + output_dir logging.debug("The onnx2tnn command:" + command + "\n") if input_names is not None: new_input_names = "" for char in input_names: if char == "[": char = ":" if char == "]": continue new_input_names += char command = command + " -input_shape " + new_input_names work_dir = "../onnx2tnn/onnx-converter/" result = cmd.run(command, work_dir=work_dir) if result == 0: logging.info("Converter ONNX to TNN model succeed!\n") else: logging.error("Converter ONNX to TNN model failed!\n") sys.exit(return_code.CONVERT_FAILED) onnx_base_name = os.path.basename(onnx_path) if align is True: if optimize is True: tnn_proto_name = onnx_base_name[:-len('.onnx' )] + '.opt' + proto_suffix tnn_model_name = onnx_base_name[:-len('.onnx' )] + '.opt' + model_suffix else: tnn_proto_name = onnx_base_name[:-len('.onnx')] + proto_suffix tnn_model_name = onnx_base_name[:-len('.onnx')] + model_suffix tnn_proto_path = os.path.join(output_dir, tnn_proto_name) tnn_model_path = os.path.join(output_dir, tnn_model_name) if input_names is None: align_model.align_model(onnx_path, tnn_proto_path, tnn_model_path, input_path, refer_path) else: align_model.align_model(onnx_path, tnn_proto_path, tnn_model_path, input_path, refer_path, new_input_names)
def print_align_message(): logging.info("{} Congratulations! {}".format("-" * 10, "-" * 10)) logging.info("The onnx model aligned with tnn model\n")
def print_not_align_message(reason): logging.info("{} Unfortunately {}" .format("-" * 10, "-" * 10)) logging.info("The onnx model not aligned with tnn model\n") logging.info("the reason " + reason + "\n") exit(-1)
def print_align_message(is_tflite=False): logging.info("{} Congratulations! {}".format("-" * 10, "-" * 10)) if is_tflite == True: logging.info("The tflite model aligned with tnn model\n") else: logging.info("The onnx model aligned with tnn model\n")
def align_model(original_model_path: str, tnn_proto_path: str, tnn_model_path: str, input_file_path: str = None, refer_path: str = None, specify_input_args: str = None, is_tflite: bool = False, debug_mode: bool = False, align_batch: bool = False) -> bool: """ 对 onnx 模型和 tnn 模型进行对齐. 当前支持模型: 单输入,单输出;单输入,多输出; :param original_model_path: :param tnn_proto_path: :param tnn_model_path: :return: """ logging.info("{} align model (tflite or ONNX vs TNN),please wait a moment {}\n" .format("-" * 10, "-" * 10)) checker.check_file_exist(tnn_proto_path) checker.check_file_exist(tnn_model_path) # list = { "input name1":{ # {"shape": [n, c,...]}, # {"data_type": 0} # }, # "input name22": { # {"shape": [n, c,...]}, # {"data_type": 0} # } # get original input info if is_tflite: original_input_info = get_input_shape_from_tflite(original_model_path) else: original_input_info = get_input_shape_from_onnx(original_model_path) # get tnn input info tnn_input_info = get_input_shape_from_tnn(tnn_proto_path) # check input if specify_input_args is not None: specify_input_info = parse_specify_input_args(specify_input_args) update_original_input_shape(original_input_info, specify_input_info) if is_tflite: check_input_lite_info(original_input_info, tnn_input_info) else: check_input_info(original_input_info, tnn_input_info) if input_file_path is None: # generate data input_path = data.gene_random_data(original_input_info) else: if os.path.exists(input_file_path): input_path = input_file_path else: logging.error("Invalid input_file_path") sys.exit(return_code.ALIGN_FAILED) if refer_path is None: if is_tflite == True: reference_output_path = run_tflite(original_model_path, input_path, original_input_info) else: reference_output_path = run_onnx(original_model_path, input_path, original_input_info) else: if os.path.exists(refer_path): reference_output_path = refer_path else: logging.error("Invalid refer_path") sys.exit(return_code.ALIGN_FAILED) logging.info("Run tnn model_check...") run_tnn_model_check(tnn_proto_path, tnn_model_path, input_path, reference_output_path, is_tflite, align_batch) if debug_mode is False: if input_file_path is None and os.path.exists(input_path): data.clean_temp_data(os.path.dirname(input_path)) if refer_path is None and os.path.exists(reference_output_path): data.clean_temp_data(reference_output_path) return True
def check_file_exist(file_path): if os.path.exists(file_path) is False: logging.info("the " + file_path + " does not exist! please make sure the file exist!\n") exit(-1)
def convert(onnx_path, output_dir=None, version="v1.0", optimize=True, half=False, align='', align_batch=False, input_path=None, refer_path=None, input_names: str = None, is_ssd=False, debug_mode: bool = False): """ 执行 onnx 转换为 tnn 的转换指令 :parameter: onnx_path: 需要转换的 onnx 文件的路径 output_path: 生成的 tnn 文件的路径 version: 转换模型的版本号 optimize: 是否需要对模型进行优化,默认是需要进行优化 half: 是否需要转为 FP16 的模型,减小模型的大小 :return return_code :exception 执行超时 """ logging.info("Converter ONNX to TNN Model...\n") checker.check_file_exist(onnx_path) try: if not is_ssd: logging.info("Converter ONNX to TNN check_onnx_dim...\n") ret, current_shape = checker.check_onnx_dim(onnx_path) logging.info("Converter ONNX to TNN check_onnx_dim...\n") if ret is False and current_shape is not None: if input_names is None: logging.info("Converter ONNX to TNN current_shape...\n") throw_exception(current_shape) if input_names is not None: input_names = input_names.strip() if ":" not in input_names and " " not in input_names: input_names = list( current_shape.keys())[0] + ":" + input_names check_input_names(input_names, current_shape) except Exception as e: print(e) logging.error( "check_onnx_dim failed, next stage of convertion may failed too\n") proto_suffix = '.tnnproto' model_suffix = '.tnnmodel' command = "python3 onnx2tnn.py " + onnx_path command = command + " -version=" + version checker.check_file_exist(onnx_path) if optimize is True: command = command + " -optimize=1" else: command = command + " -optimize=0" if half is True: command = command + " -half=1" else: command = command + " -half=0" if output_dir is None: output_dir = os.path.dirname(onnx_path) checker.check_file_exist(output_dir) command = command + " -o " + output_dir if input_names is not None: command = command + " -input_shape " + input_names logging.debug("The onnx2tnn command:" + command + "\n") work_dir = "../onnx2tnn/onnx-converter/" result = cmd.run(command, work_dir=work_dir) if result == 0: logging.info("Converter ONNX to TNN model succeed!\n") else: logging.error("Converter ONNX to TNN model failed!\n") sys.exit(return_code.CONVERT_FAILED) onnx_base_name = os.path.basename(onnx_path) if optimize is True: tnn_proto_name = onnx_base_name[:-len('.onnx')] + '.opt' + proto_suffix tnn_model_name = onnx_base_name[:-len('.onnx')] + '.opt' + model_suffix else: tnn_proto_name = onnx_base_name[:-len('.onnx')] + proto_suffix tnn_model_name = onnx_base_name[:-len('.onnx')] + model_suffix tnn_proto_path = os.path.join(output_dir, tnn_proto_name) tnn_model_path = os.path.join(output_dir, tnn_model_name) if align == 'output' or align_batch is True: if input_names is None: align_model.align_model(onnx_path, tnn_proto_path, tnn_model_path, input_path, refer_path, debug_mode=debug_mode, align_batch=align_batch) else: align_model.align_model(onnx_path, tnn_proto_path, tnn_model_path, input_path, refer_path, input_names, debug_mode=debug_mode, align_batch=align_batch) elif align == 'all': is_align_all = (align == 'all') align_model.align_all(onnx_path, tnn_proto_path, is_align_all, input_names, input_path, refer_path)