69 lines
2.6 KiB
Python
69 lines
2.6 KiB
Python
#!/usr/bin/env python3
|
|
import argparse
|
|
import sys
|
|
import os
|
|
import logging
|
|
from modules import opt_module
|
|
from utils import log
|
|
from utils import util
|
|
|
|
|
|
def arg_parse():
|
|
desc = 'Environment variables:\n - COMPILER_BIN_DIR: specify compiler bin dir\n\n'
|
|
desc += 'Sub-modules:\n - image_cut_search: 720|530|730|630|540\n'
|
|
parser = argparse.ArgumentParser(description=desc, formatter_class=argparse.RawTextHelpFormatter)
|
|
parser.add_argument("platform", help="HW platform <520|720|530|730|630|540>")
|
|
parser.add_argument("model", help="onnx or bie model path")
|
|
parser.add_argument("config", help="compiler base config. config.json")
|
|
parser.add_argument('output_dir', help='log output directory')
|
|
parser.add_argument("-d", "--debug", action="store_true", default=False, help="debug mode. Show more debug message")
|
|
|
|
return parser.parse_args()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
|
# get opt wrapper directory
|
|
opt_dir = os.path.dirname(os.path.abspath(__file__))
|
|
|
|
# get arguments
|
|
args = arg_parse()
|
|
platform = args.platform
|
|
model = args.model
|
|
base_config = args.config
|
|
output_dir = args.output_dir
|
|
log_level = logging.DEBUG if args.debug else logging.INFO
|
|
|
|
# create logger and output directory
|
|
log_path = output_dir if output_dir else '__magic_tmp__'
|
|
remove_log = not output_dir
|
|
log = log.Log(log_path, remove_log)
|
|
log.copy_files(model, base_config)
|
|
log.logger.setLevel(log_level)
|
|
|
|
if output_dir and not util.is_file_name_valid(output_dir):
|
|
log.logger.error('Please give a valid -o/OUTPUT_DIR. "{}" is invalid'.format(output_dir))
|
|
sys.exit(-1)
|
|
|
|
# set compiler binary path to environment. the env is used in sub-module
|
|
compiler_bin = os.environ['COMPILER_BIN'] = util.get_compiler_bin(opt_dir)
|
|
if not util.check_file_exist(compiler_bin):
|
|
log.logger.error('compiler binary [path: {}] not exist!'.format(compiler_bin))
|
|
sys.exit(-1)
|
|
|
|
|
|
# run optimize modules
|
|
opt_modules = opt_module.OptModules(platform, model, base_config, opt_dir, log)
|
|
log.logger.info('==== Information ====')
|
|
log.logger.info('opt_compile : {}'.format(' '.join(sys.argv)))
|
|
log.logger.info('platform : {}'.format(platform))
|
|
log.logger.info('model : {}'.format(model))
|
|
log.logger.info('base config : {}'.format(base_config))
|
|
log.logger.info('sub-modules : {}'.format(opt_modules.modules))
|
|
log.logger.info('log directory : {}'.format(log_path))
|
|
log.logger.info('compiler bin : {}'.format(compiler_bin))
|
|
log.logger.info('=====================')
|
|
opt_modules.run()
|
|
|
|
# TODO: find best order for opt_modules
|