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))