| |
| import argparse |
| import os |
| import os.path as osp |
|
|
|
|
| def parse_args(): |
| parser = argparse.ArgumentParser( |
| description='Convert benchmark model json to script') |
| parser.add_argument( |
| 'txt_path', type=str, help='txt path output by benchmark_filter') |
| 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 determine_gpus(cfg_name): |
| gpus = 8 |
| gpus_pre_node = 8 |
|
|
| if cfg_name.find('16x') >= 0: |
| gpus = 16 |
| elif cfg_name.find('4xb4') >= 0: |
| gpus = 4 |
| gpus_pre_node = 4 |
| elif 'lad' in cfg_name: |
| gpus = 2 |
| gpus_pre_node = 2 |
|
|
| return gpus, gpus_pre_node |
|
|
|
|
| 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"') |
|
|
| root_name = './tools' |
| train_script_name = osp.join(root_name, 'slurm_train.sh') |
|
|
| commands = [] |
| partition_name = 'PARTITION=$1 ' |
| commands.append(partition_name) |
| commands.append('\n') |
|
|
| work_dir = 'WORK_DIR=$2 ' |
| commands.append(work_dir) |
| commands.append('\n') |
|
|
| cpus_pre_task = 'CPUS_PER_TASK=${3:-4} ' |
| commands.append(cpus_pre_task) |
| commands.append('\n') |
| commands.append('\n') |
|
|
| with open(args.txt_path, 'r') as f: |
| model_cfgs = f.readlines() |
| for i, cfg in enumerate(model_cfgs): |
| cfg = cfg.strip() |
| if len(cfg) == 0: |
| continue |
| |
| echo_info = f'echo \'{cfg}\' &' |
| commands.append(echo_info) |
| commands.append('\n') |
|
|
| fname, _ = osp.splitext(osp.basename(cfg)) |
| out_fname = '$WORK_DIR/' + fname |
|
|
| gpus, gpus_pre_node = determine_gpus(cfg) |
| command_info = f'GPUS={gpus} GPUS_PER_NODE={gpus_pre_node} ' \ |
| f'CPUS_PER_TASK=$CPUS_PRE_TASK {train_script_name} ' |
| command_info += '$PARTITION ' |
| command_info += f'{fname} ' |
| command_info += f'{cfg} ' |
| command_info += f'{out_fname} ' |
|
|
| command_info += '--cfg-options default_hooks.checkpoint.' \ |
| 'max_keep_ckpts=1 ' |
| command_info += '&' |
|
|
| commands.append(command_info) |
|
|
| if i < len(model_cfgs): |
| commands.append('\n') |
|
|
| 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() |
|
|
