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(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 convert(tf_path, input_names, output_names, output_dir, version, optimize, half, align=False, not_fold_const=False, input_path=None, refer_path=None, debug: bool = False, debug_mode: bool = False): 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 = time.strftime('%Y%m%d%H%M', time.localtime()) 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, debug_mode=debug_mode) onnx_base_name = os.path.basename(onnx_path) tnn_proto_name = onnx_base_name[:-len('.onnx')] + ( '.opt.tnnproto' if optimize else ".tnnproto") tnn_proto_path = os.path.join(output_dir, tnn_proto_name) fix_tnn_output.fix_tnn_output(tnn_proto_path)
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, is_ssd=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) if not is_ssd: 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 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 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 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 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)