|
import argparse |
|
import re |
|
from os import path as osp |
|
|
|
|
|
def parse_args(): |
|
parser = argparse.ArgumentParser( |
|
description='Generate benchmark training/testing scripts') |
|
parser.add_argument( |
|
'--input_file', |
|
required=False, |
|
type=str, |
|
help='Input file containing the paths ' |
|
'of configs to be trained/tested.') |
|
parser.add_argument( |
|
'--output_file', |
|
required=True, |
|
type=str, |
|
help='Output file containing the ' |
|
'commands to train/test selected models.') |
|
parser.add_argument( |
|
'--gpus_per_node', |
|
type=int, |
|
default=8, |
|
help='GPUs per node config for slurm, ' |
|
'should be set according to your slurm environment') |
|
parser.add_argument( |
|
'--cpus_per_task', |
|
type=int, |
|
default=5, |
|
help='CPUs per task config for slurm, ' |
|
'should be set according to your slurm environment') |
|
parser.add_argument( |
|
'--gpus', |
|
type=int, |
|
default=8, |
|
help='Totally used num of GPUs config for slurm (in testing), ' |
|
'should be set according to your slurm environment') |
|
parser.add_argument( |
|
'--mode', type=str, default='train', help='Train or test') |
|
parser.add_argument( |
|
'--long_work_dir', |
|
action='store_true', |
|
help='Whether use full relative path of config as work dir') |
|
parser.add_argument( |
|
'--max_keep_ckpts', |
|
type=int, |
|
default=1, |
|
help='The max number of checkpoints saved in training') |
|
parser.add_argument( |
|
'--full_log', |
|
action='store_true', |
|
help='Whether save full log in a file') |
|
|
|
args = parser.parse_args() |
|
return args |
|
|
|
|
|
args = parse_args() |
|
assert args.mode in ['train', 'test'], 'Currently we only support ' \ |
|
'automatically generating training or testing scripts.' |
|
|
|
config_paths = [] |
|
|
|
if args.input_file is not None: |
|
with open(args.input_file, 'r') as fi: |
|
config_paths = fi.read().strip().split('\n') |
|
else: |
|
while True: |
|
print('Please type a config path and ' |
|
'press enter (press enter directly to exit):') |
|
config_path = input() |
|
if config_path != '': |
|
config_paths.append(config_path) |
|
else: |
|
break |
|
|
|
script = '''PARTITION=$1 |
|
CHECKPOINT_DIR=$2 |
|
|
|
''' |
|
|
|
if args.mode == 'train': |
|
for i, config_path in enumerate(config_paths): |
|
root_dir = osp.dirname(osp.dirname(osp.abspath(__file__))) |
|
if not osp.exists(osp.join(root_dir, config_path)): |
|
print(f'Invalid config path (does not exist):\n{config_path}') |
|
continue |
|
|
|
config_name = config_path.split('/')[-1][:-3] |
|
match_obj = re.match(r'^.*_[0-9]+x([0-9]+)_.*$', config_name) |
|
if match_obj is None: |
|
print(f'Invalid config path (no GPU num in ' |
|
f'config name):\n{config_path}') |
|
continue |
|
|
|
gpu_num = int(match_obj.group(1)) |
|
work_dir_name = config_path if args.long_work_dir else config_name |
|
|
|
script += f"echo '{config_path}' &\n" |
|
if args.full_log: |
|
script += f'mkdir -p $CHECKPOINT_DIR/{work_dir_name}\n' |
|
|
|
|
|
script += f'GPUS={gpu_num} GPUS_PER_NODE={args.gpus_per_node} ' \ |
|
f'CPUS_PER_TASK={args.cpus_per_task} ' \ |
|
f'./tools/slurm_train.sh $PARTITION {config_name} ' \ |
|
f'{config_path} \\\n' |
|
script += f'$CHECKPOINT_DIR/{work_dir_name} --cfg-options ' \ |
|
f'checkpoint_config.max_keep_ckpts=' \ |
|
f'{args.max_keep_ckpts} \\\n' \ |
|
|
|
|
|
|
|
if args.full_log: |
|
script += f'2>&1|tee $CHECKPOINT_DIR/{work_dir_name}' \ |
|
f'/FULL_LOG.txt &\n' |
|
else: |
|
script += '>/dev/null &\n' |
|
|
|
if i != len(config_paths) - 1: |
|
script += '\n' |
|
|
|
print(f'Successfully generated script for {config_name}') |
|
|
|
with open(args.output_file, 'w') as fo: |
|
fo.write(script) |
|
|
|
elif args.mode == 'test': |
|
for i, config_path in enumerate(config_paths): |
|
root_dir = osp.dirname(osp.dirname(osp.abspath(__file__))) |
|
if not osp.exists(osp.join(root_dir, config_path)): |
|
print(f'Invalid config path (does not exist):\n{config_path}') |
|
continue |
|
|
|
config_name = config_path.split('/')[-1][:-3] |
|
|
|
tasks = { |
|
'scannet_seg', 'scannet', 's3dis_seg', 'sunrgbd', 'kitti', 'nus', |
|
'lyft', 'waymo' |
|
} |
|
eval_option = None |
|
for task in tasks: |
|
if task in config_name: |
|
eval_option = task |
|
break |
|
if eval_option is None: |
|
print(f'Invalid config path (invalid task):\n{config_path}') |
|
continue |
|
|
|
work_dir_name = config_path if args.long_work_dir else config_name |
|
|
|
script += f"echo '{config_path}' &\n" |
|
if args.full_log: |
|
script += f'mkdir -p $CHECKPOINT_DIR/{work_dir_name}\n' |
|
|
|
|
|
script += f'GPUS={args.gpus} GPUS_PER_NODE={args.gpus_per_node} ' \ |
|
f'CPUS_PER_TASK={args.cpus_per_task} ' \ |
|
f'./tools/slurm_test.sh $PARTITION {config_name} ' \ |
|
f'{config_path} \\\n' |
|
script += f'$CHECKPOINT_DIR/{work_dir_name}/latest.pth ' \ |
|
|
|
if eval_option in ['scannet_seg', 's3dis_seg']: |
|
script += '--eval mIoU \\\n' |
|
elif eval_option in ['scannet', 'sunrgbd', 'kitti', 'nus']: |
|
script += '--eval map \\\n' |
|
elif eval_option in ['lyft']: |
|
script += f'--format-only --eval-options jsonfile_prefix=' \ |
|
f'$CHECKPOINT_DIR/{work_dir_name}/results_challenge ' \ |
|
f'csv_savepath=$CHECKPOINT_DIR/{work_dir_name}/' \ |
|
f'results_challenge.csv \\\n' |
|
elif eval_option in ['waymo']: |
|
script += f'--eval waymo --eval-options pklfile_prefix=' \ |
|
f'$CHECKPOINT_DIR/{work_dir_name}/kitti_results ' \ |
|
f'submission_prefix=$CHECKPOINT_DIR/{work_dir_name}/' \ |
|
f'kitti_results \\\n' |
|
|
|
|
|
|
|
if args.full_log: |
|
script += f'2>&1|tee $CHECKPOINT_DIR/{work_dir_name}' \ |
|
f'/FULL_LOG.txt &\n' |
|
else: |
|
script += '>/dev/null &\n' |
|
|
|
if i != len(config_paths) - 1: |
|
script += '\n' |
|
|
|
print(f'Successfully generated script for {config_name}') |
|
|
|
with open(args.output_file, 'w') as fo: |
|
fo.write(script) |
|
|
