def main(args, exp_config, train_set, val_set, test_set): # Record settings exp_config.update({ 'model': args['model'], 'in_feats': args['node_featurizer'].feat_size(), 'n_tasks': args['n_tasks'] }) # Set up directory for saving results args = init_trial_path(args) train_loader = DataLoader(dataset=train_set, batch_size=exp_config['batch_size'], shuffle=True, collate_fn=collate_molgraphs) val_loader = DataLoader(dataset=val_set, batch_size=exp_config['batch_size'], collate_fn=collate_molgraphs) test_loader = DataLoader(dataset=test_set, batch_size=exp_config['batch_size'], collate_fn=collate_molgraphs) model = load_model(exp_config).to(args['device']) loss_criterion = nn.SmoothL1Loss(reduction='none') optimizer = Adam(model.parameters(), lr=exp_config['lr'], weight_decay=exp_config['weight_decay']) stopper = EarlyStopping(patience=exp_config['patience'], filename=args['trial_path'] + '/model.pth') for epoch in range(args['num_epochs']): # Train run_a_train_epoch(args, epoch, model, train_loader, loss_criterion, optimizer) # Validation and early stop val_score = run_an_eval_epoch(args, model, val_loader) early_stop = stopper.step(val_score, model) print( 'epoch {:d}/{:d}, validation {} {:.4f}, best validation {} {:.4f}'. format(epoch + 1, args['num_epochs'], args['metric'], val_score, args['metric'], stopper.best_score)) if early_stop: break stopper.load_checkpoint(model) test_score = run_an_eval_epoch(args, model, test_loader) print('test {} {:.4f}'.format(args['metric'], test_score)) with open(args['trial_path'] + '/eval.txt', 'w') as f: f.write('Best val {}: {}\n'.format(args['metric'], stopper.best_score)) f.write('Test {}: {}\n'.format(args['metric'], test_score)) with open(args['trial_path'] + '/configure.json', 'w') as f: json.dump(exp_config, f, indent=2) return args['trial_path'], stopper.best_score
def objective(hyperparams): configure = deepcopy(args) save_path = init_trial_path(args) val_metrics, test_metrics = main(save_path, configure, hyperparams) if args['metric'] in ['roc_auc']: # To maximize a non-negative value is equivalent to minimize its opposite number val_metric_to_minimize = -1 * np.mean(val_metrics) else: val_metric_to_minimize = np.mean(val_metrics) results.append( (save_path, val_metric_to_minimize, val_metrics, test_metrics)) return val_metric_to_minimize
def main(args, train_set, val_set, test_set): # Set up directory for saving results args = init_trial_path(args) train_loader = DataLoader(dataset=train_set, batch_size=args['batch_size'], shuffle=True, collate_fn=collate_molgraphs) val_loader = DataLoader(dataset=val_set, batch_size=args['batch_size'], collate_fn=collate_molgraphs) test_loader = DataLoader(dataset=test_set, batch_size=args['batch_size'], collate_fn=collate_molgraphs) model = load_model(args).to(args['device']) loss_criterion = nn.BCEWithLogitsLoss(reduction='none') optimizer = Adam(model.parameters(), lr=args['lr'], weight_decay=args['weight_decay']) stopper = EarlyStopping(patience=args['patience'], filename=args['trial_path'] + '/model.pth') for epoch in range(args['num_epochs']): # Train run_a_train_epoch(args, epoch, model, train_loader, loss_criterion, optimizer) # Validation and early stop val_score = run_an_eval_epoch(args, model, val_loader) early_stop = stopper.step(val_score, model) print( 'epoch {:d}/{:d}, validation {} {:.4f}, best validation {} {:.4f}'. format(epoch + 1, args['num_epochs'], args['metric'], val_score, args['metric'], stopper.best_score)) if early_stop: break test_score = run_an_eval_epoch(args, model, test_loader) print('test {} {:.4f}'.format(args['metric'], test_score)) with open(args['trial_path'] + '/eval.txt', 'w') as f: f.write('Best val {}: {}\n'.format(args['metric'], stopper.best_score)) f.write('Test {}: {}\n'.format(args['metric'], test_score)) return args, stopper.best_score