67 lines
1.9 KiB
Python
67 lines
1.9 KiB
Python
from utils.common import *
|
|
from utils.helper import *
|
|
from utils.load_config import BatchConfig
|
|
from utils.generate_batch_config import generate_batch_bconf, generate_batch_conf
|
|
|
|
import argparse
|
|
import os
|
|
import shutil
|
|
import subprocess
|
|
|
|
# Parse arguments
|
|
argparser = argparse.ArgumentParser(
|
|
description="Run batch compiler for KDP 720")
|
|
|
|
argparser.add_argument(
|
|
'-c',
|
|
'--config',
|
|
type=str,
|
|
default='/data1/batch_input_params.json',
|
|
help="path of input_params.json")
|
|
|
|
argparser.add_argument(
|
|
'--flatbuffer',
|
|
action='store_true',
|
|
default=False,
|
|
help="enable new flatbuffer mode")
|
|
|
|
args = argparser.parse_args()
|
|
|
|
# Load config
|
|
batch_config = BatchConfig(args.config)
|
|
prepare_tmp_folder()
|
|
output_path = prepare_result_folder('batch_compile')
|
|
|
|
# check config
|
|
bie_models = dict()
|
|
for batch_model in batch_config.model_list:
|
|
if batch_model.path.endswith(".bie"):
|
|
continue
|
|
elif batch_model.path.endswith(".onnx") and batch_model.radix is not None:
|
|
continue
|
|
else:
|
|
raise ValueError("Currently, batch compile only support models after fix point analysis.")
|
|
|
|
# Generate batch compile config
|
|
generate_batch_conf(batch_config.encryption,
|
|
'720',
|
|
output_path + '/batch_compile_config.json',
|
|
weight_compress=batch_config.weight_compress,
|
|
flatbuffer=args.flatbuffer)
|
|
|
|
# Generate batch compile b_conf
|
|
generate_batch_bconf(batch_config,
|
|
output_path + '/batch_compile_config.json',
|
|
output_path + '/batch_compile_bconfig.json')
|
|
|
|
# Run batch compile
|
|
os.chdir(output_path)
|
|
commands = [LIBS_FOLDER + '/compiler/batch_compile',
|
|
'-T', '720',
|
|
output_path + '/batch_compile_bconfig.json',
|
|
'-s', '-t', get_toolchain_version()]
|
|
if batch_config.dedicated_buffer:
|
|
commands.append('-o')
|
|
subprocess.run(commands, check=True)
|
|
|
|
delete_tmp_folder() |