|
import argparse |
|
import os |
|
|
|
import yaml |
|
import torch |
|
import torch.multiprocessing as mp |
|
|
|
import utils |
|
from utils.experiment import * |
|
from trainer import Trainer |
|
|
|
|
|
def parse_args(): |
|
parser = argparse.ArgumentParser() |
|
parser.add_argument('--cfg') |
|
parser.add_argument('--load-root', default='data') |
|
parser.add_argument('--save-root', default='save') |
|
parser.add_argument('--name', '-n', default=None) |
|
parser.add_argument('--tag', default=None) |
|
parser.add_argument('--cudnn', action='store_true') |
|
parser.add_argument('--port-offset', '-p', type=int, default=0) |
|
parser.add_argument('--wandb-upload', '-w', action='store_true') |
|
args = parser.parse_args() |
|
|
|
return args |
|
|
|
|
|
def make_cfg(args): |
|
with open(args.cfg, 'r') as f: |
|
cfg = yaml.load(f, Loader=yaml.FullLoader) |
|
|
|
def translate_cfg_(d): |
|
for k, v in d.items(): |
|
if isinstance(v, dict): |
|
translate_cfg_(v) |
|
elif isinstance(v, str): |
|
d[k] = v.replace('$load_root$', args.load_root) |
|
translate_cfg_(cfg) |
|
|
|
if args.name is None: |
|
exp_name = os.path.basename(args.cfg).split('.')[0].replace('_benchmark', '').replace('_demo', '') |
|
else: |
|
exp_name = args.name |
|
if args.tag is not None: |
|
exp_name += '_' + args.tag |
|
|
|
env = dict() |
|
env['exp_name'] = exp_name + '_' + cfg['exp_name'] |
|
env['save_dir'] = os.path.join(args.save_root, env['exp_name']) |
|
|
|
env['cudnn'] = args.cudnn |
|
env['port'] = str(29600 + args.port_offset) |
|
env['wandb_upload'] = args.wandb_upload |
|
cfg['env'] = env |
|
|
|
return cfg |
|
|
|
|
|
def main(): |
|
args = parse_args() |
|
|
|
cfgs = make_cfg(args) |
|
|
|
init_experiment(cfgs) |
|
init_distributed_mode(cfgs) |
|
print('here') |
|
init_deterministic(cfgs['seed']) |
|
|
|
trainer = Trainer(cfgs) |
|
|
|
if cfgs['mode'] == 'train': |
|
trainer.train() |
|
elif cfgs['mode'] == 'validate': |
|
trainer.validate() |
|
elif cfgs['mode'] == 'test': |
|
trainer.test() |
|
elif cfgs['mode'] == 'demo': |
|
trainer.demo() |
|
|
|
|
|
|
|
if __name__ == '__main__': |
|
main() |
|
|
