|
- # Copyright (c) OpenMMLab. All rights reserved.
- 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(
- '--partition',
- type=str,
- default='openmmlab',
- help='slurm partition name')
- parser.add_argument(
- '--max-keep-ckpts',
- type=int,
- default=1,
- help='The maximum checkpoints to keep')
- 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 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"')
-
- partition = args.partition # cluster name
-
- root_name = './tools'
- train_script_name = osp.join(root_name, 'slurm_train.sh')
- # stdout is no output
- stdout_cfg = '>/dev/null'
-
- max_keep_ckpts = args.max_keep_ckpts
-
- commands = []
- 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
- # print cfg name
- echo_info = f'echo \'{cfg}\' &'
- commands.append(echo_info)
- commands.append('\n')
-
- fname, _ = osp.splitext(osp.basename(cfg))
- out_fname = osp.join(root_name, 'work_dir', fname)
- # default setting
- if cfg.find('16x') >= 0:
- command_info = f'GPUS=16 GPUS_PER_NODE=8 ' \
- f'CPUS_PER_TASK=2 {train_script_name} '
- elif cfg.find('gn-head_4x4_1x_coco.py') >= 0 or \
- cfg.find('gn-head_4x4_2x_coco.py') >= 0:
- command_info = f'GPUS=4 GPUS_PER_NODE=4 ' \
- f'CPUS_PER_TASK=2 {train_script_name} '
- else:
- command_info = f'GPUS=8 GPUS_PER_NODE=8 ' \
- f'CPUS_PER_TASK=2 {train_script_name} '
- command_info += f'{partition} '
- command_info += f'{fname} '
- command_info += f'{cfg} '
- command_info += f'{out_fname} '
- if max_keep_ckpts:
- command_info += f'--cfg-options ' \
- f'checkpoint_config.max_keep_ckpts=' \
- f'{max_keep_ckpts}' + ' '
- command_info += f'{stdout_cfg} &'
-
- 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()
|