76 lines
3.1 KiB
Python
76 lines
3.1 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: 730|1140\n'
|
|
parser = argparse.ArgumentParser(description=desc, formatter_class=argparse.RawTextHelpFormatter)
|
|
parser.add_argument("platform", help="HW platform <730|1140>")
|
|
parser.add_argument("model", help="onnx or bie model path")
|
|
parser.add_argument("config", help="compiler base config. config.json")
|
|
parser.add_argument('log_dir', help='directory to store debug files')
|
|
parser.add_argument("output_dir", help="directory to store exported files")
|
|
parser.add_argument("--image_cut_search_args", type=str, help="arguments for image cut search")
|
|
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
|
|
log_dir = args.log_dir
|
|
image_cut_search_args = args.image_cut_search_args
|
|
output_dir = args.output_dir
|
|
opt_compile_cmd = f'{sys.argv[0]} {platform} {model} {base_config} {log_dir} {output_dir} --image_cut_search_args "{image_cut_search_args}"'
|
|
|
|
# create logger and output directory
|
|
log_level = logging.DEBUG if args.debug else logging.INFO
|
|
log = log.Log(log_dir, output_dir)
|
|
log.copy_files(model, base_config)
|
|
log.logger.setLevel(log_level)
|
|
|
|
if log_dir and not util.is_file_name_valid(log_dir):
|
|
log.logger.error('Please give a valid LOG_DIR. "{}" is invalid'.format(log_dir))
|
|
sys.exit(-1)
|
|
|
|
if output_dir and not util.is_file_name_valid(output_dir):
|
|
log.logger.error('Please give a valid 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, image_cut_search_args)
|
|
log.logger.info('==== Information ====')
|
|
log.logger.info('opt_compile : {}'.format(opt_compile_cmd))
|
|
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 dir : {}'.format(log_dir))
|
|
log.logger.info('output dir : {}'.format(output_dir))
|
|
log.logger.info('compiler bin : {}'.format(compiler_bin))
|
|
log.logger.info('=====================')
|
|
opt_modules.run()
|
|
|
|
# TODO: find best order for opt_modules
|