args = parse_args()
# Load config file
load_cfg(cfg, args)
set_out_dir(cfg.out_dir, args.cfg_file)
# Set Pytorch environment
torch.set_num_threads(cfg.num_threads)
dump_cfg(cfg)
# Repeat for different random seeds
for i in range(args.repeat):
set_run_dir(cfg.out_dir)
set_printing()
# Set configurations for each run
cfg.seed = cfg.seed + 1
seed_everything(cfg.seed)
auto_select_device()
# Set machine learning pipeline
datamodule = GraphGymDataModule()
model = create_model()
# Print model info
logging.info(model)
logging.info(cfg)
cfg.params = params_count(model)
logging.info('Num parameters: %s', cfg.params)
train(model, datamodule, logger=True)
# Aggregate results from different seeds
agg_runs(cfg.out_dir, cfg.metric_best)
# When being launched in batch mode, mark a yaml as done
if args.mark_done:
os.rename(args.cfg_file, f'{args.cfg_file}_done')
set_printing()
# Repeat for different random seeds
for i in range(args.repeat):
# Set configurations for each run
seed_everything(cfg.seed + i)
auto_select_device()
set_run_dir(cfg.out_dir, args.cfg_file)
# Set machine learning pipeline
loaders = create_loader()
loggers = create_logger()
model = create_model()
optimizer = create_optimizer(model.parameters(),
new_optimizer_config(cfg))
scheduler = create_scheduler(optimizer, new_scheduler_config(cfg))
# Print model info
logging.info(model)
logging.info(cfg)
cfg.params = params_count(model)
logging.info('Num parameters: {}'.format(cfg.params))
# Start training
if cfg.train.mode == 'standard':
train(loggers, loaders, model, optimizer, scheduler)
else:
train_dict[cfg.train.mode](loggers, loaders, model, optimizer,
scheduler)
# Aggregate results from different seeds
agg_runs(set_agg_dir(cfg.out_dir, args.cfg_file), cfg.metric_best)
# When being launched in batch mode, mark a yaml as done
if args.mark_done:
os.rename(args.cfg_file, '{}_done'.format(args.cfg_file))