| | |
| | import argparse |
| | import os |
| | import os.path as osp |
| |
|
| | from mmengine import Config |
| |
|
| |
|
| | def parse_args(): |
| | parser = argparse.ArgumentParser( |
| | description='Convert benchmark model list to script') |
| | parser.add_argument('config', help='test config file path') |
| | parser.add_argument('--port', type=int, default=29666, help='dist port') |
| | parser.add_argument( |
| | '--run', action='store_true', help='run script directly') |
| | parser.add_argument( |
| | '--out', type=str, help='path to save model benchmark script') |
| |
|
| | args = parser.parse_args() |
| | return args |
| |
|
| |
|
| | def process_model_info(model_info, work_dir): |
| | config = model_info['config'].strip() |
| | fname, _ = osp.splitext(osp.basename(config)) |
| | job_name = fname |
| | work_dir = '$WORK_DIR/' + fname |
| | checkpoint = model_info['checkpoint'].strip() |
| | return dict( |
| | config=config, |
| | job_name=job_name, |
| | work_dir=work_dir, |
| | checkpoint=checkpoint) |
| |
|
| |
|
| | def create_test_bash_info(commands, model_test_dict, port, script_name, |
| | partition): |
| | config = model_test_dict['config'] |
| | job_name = model_test_dict['job_name'] |
| | checkpoint = model_test_dict['checkpoint'] |
| | work_dir = model_test_dict['work_dir'] |
| |
|
| | echo_info = f' \necho \'{config}\' &' |
| | commands.append(echo_info) |
| | commands.append('\n') |
| |
|
| | command_info = f'GPUS=8 GPUS_PER_NODE=8 ' \ |
| | f'CPUS_PER_TASK=$CPUS_PRE_TASK {script_name} ' |
| |
|
| | command_info += f'{partition} ' |
| | command_info += f'{job_name} ' |
| | command_info += f'{config} ' |
| | command_info += f'$CHECKPOINT_DIR/{checkpoint} ' |
| | command_info += f'--work-dir {work_dir} ' |
| |
|
| | command_info += f'--cfg-option env_cfg.dist_cfg.port={port} ' |
| | command_info += ' &' |
| |
|
| | commands.append(command_info) |
| |
|
| |
|
| | def main(): |
| | args = parse_args() |
| | if args.out: |
| | out_suffix = args.out.split('.')[-1] |
| | assert args.out.endswith('.sh'), \ |
| | f'Expected out file path suffix is .sh, but get .{out_suffix}' |
| | assert args.out or args.run, \ |
| | ('Please specify at least one operation (save/run/ the ' |
| | 'script) with the argument "--out" or "--run"') |
| |
|
| | commands = [] |
| | partition_name = 'PARTITION=$1 ' |
| | commands.append(partition_name) |
| | commands.append('\n') |
| |
|
| | checkpoint_root = 'CHECKPOINT_DIR=$2 ' |
| | commands.append(checkpoint_root) |
| | commands.append('\n') |
| |
|
| | work_dir = 'WORK_DIR=$3 ' |
| | commands.append(work_dir) |
| | commands.append('\n') |
| |
|
| | cpus_pre_task = 'CPUS_PER_TASK=${4:-2} ' |
| | commands.append(cpus_pre_task) |
| | commands.append('\n') |
| |
|
| | script_name = osp.join('tools', 'slurm_test.sh') |
| | port = args.port |
| |
|
| | cfg = Config.fromfile(args.config) |
| |
|
| | for model_key in cfg: |
| | model_infos = cfg[model_key] |
| | if not isinstance(model_infos, list): |
| | model_infos = [model_infos] |
| | for model_info in model_infos: |
| | print('processing: ', model_info['config']) |
| | model_test_dict = process_model_info(model_info, work_dir) |
| | create_test_bash_info(commands, model_test_dict, port, script_name, |
| | '$PARTITION') |
| | port += 1 |
| |
|
| | command_str = ''.join(commands) |
| | if args.out: |
| | with open(args.out, 'w') as f: |
| | f.write(command_str) |
| | if args.run: |
| | os.system(command_str) |
| |
|
| |
|
| | if __name__ == '__main__': |
| | main() |
| |
|
