|
|
|
|
|
import argparse |
|
|
import os.path as osp |
|
|
|
|
|
|
|
|
config_8gpu_list = [ |
|
|
'configs/swin/upernet_swin_tiny_patch4_window7_512x512_160k_ade20k_pretrain_224x224_1K.py', |
|
|
'configs/vit/upernet_vit-b16_ln_mln_512x512_160k_ade20k.py', |
|
|
'configs/vit/upernet_deit-s16_ln_mln_512x512_160k_ade20k.py', |
|
|
] |
|
|
|
|
|
|
|
|
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('--port', type=int, default=24727, help='dist port') |
|
|
parser.add_argument( |
|
|
'--out', |
|
|
type=str, |
|
|
default='.dev_scripts/benchmark_train.sh', |
|
|
help='path to save model benchmark script') |
|
|
|
|
|
args = parser.parse_args() |
|
|
return args |
|
|
|
|
|
|
|
|
def create_train_bash_info(commands, config, script_name, partition, port): |
|
|
cfg = config.strip() |
|
|
|
|
|
|
|
|
echo_info = f'echo \'{cfg}\' &' |
|
|
commands.append(echo_info) |
|
|
commands.append('\n') |
|
|
|
|
|
_, model_name = osp.split(osp.dirname(cfg)) |
|
|
config_name, _ = osp.splitext(osp.basename(cfg)) |
|
|
|
|
|
if cfg in config_8gpu_list: |
|
|
command_info = f'GPUS=8 GPUS_PER_NODE=8 ' \ |
|
|
f'CPUS_PER_TASK=2 {script_name} ' |
|
|
else: |
|
|
command_info = f'GPUS=4 GPUS_PER_NODE=4 ' \ |
|
|
f'CPUS_PER_TASK=2 {script_name} ' |
|
|
command_info += f'{partition} ' |
|
|
command_info += f'{config_name} ' |
|
|
command_info += f'{cfg} ' |
|
|
command_info += f'--cfg-options ' \ |
|
|
f'checkpoint_config.max_keep_ckpts=1 ' \ |
|
|
f'dist_params.port={port} ' |
|
|
command_info += f'--work-dir work_dirs/{model_name}/{config_name} ' |
|
|
|
|
|
command_info += '>/dev/null &' |
|
|
|
|
|
commands.append(command_info) |
|
|
commands.append('\n') |
|
|
|
|
|
|
|
|
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}' |
|
|
|
|
|
root_name = './tools' |
|
|
script_name = osp.join(root_name, 'slurm_train.sh') |
|
|
port = args.port |
|
|
partition_name = 'PARTITION=$1' |
|
|
|
|
|
commands = [] |
|
|
commands.append(partition_name) |
|
|
commands.append('\n') |
|
|
commands.append('\n') |
|
|
|
|
|
with open(args.txt_path) as f: |
|
|
model_cfgs = f.readlines() |
|
|
for i, cfg in enumerate(model_cfgs): |
|
|
create_train_bash_info(commands, cfg, script_name, '$PARTITION', |
|
|
port) |
|
|
port += 1 |
|
|
|
|
|
command_str = ''.join(commands) |
|
|
if args.out: |
|
|
with open(args.out, 'w') as f: |
|
|
f.write(command_str) |
|
|
|
|
|
|
|
|
if __name__ == '__main__': |
|
|
main() |
|
|
|
