def makeTrainer(*, task='h**o', device='cuda', lr=3e-3, bs=75, num_epochs=500,network=MolecLieResNet, net_config={'k':1536,'nbhd':100,'act':'swish','group':lieGroups.T(3), 'bn':True,'aug':True,'mean':True,'num_layers':6}, recenter=False, subsample=False, trainer_config={'log_dir':None,'log_suffix':''}):#,'log_args':{'timeFrac':1/4,'minPeriod':0}}): # Create Training set and model device = torch.device(device) with FixedNumpySeed(0): datasets, num_species, charge_scale = QM9datasets() if subsample: datasets.update(split_dataset(datasets['train'],{'train':subsample})) ds_stats = datasets['train'].stats[task] if recenter: m = datasets['train'].data['charges']>0 pos = datasets['train'].data['positions'][m] mean,std = pos.mean(dim=0),1#pos.std() for ds in datasets.values(): ds.data['positions'] = (ds.data['positions']-mean[None,None,:])/std model = network(num_species,charge_scale,**net_config).to(device) # Create train and Val(Test) dataloaders and move elems to gpu dataloaders = {key:LoaderTo(DataLoader(dataset,batch_size=bs,num_workers=0, shuffle=(key=='train'),pin_memory=False,collate_fn=collate_fn,drop_last=True), device) for key,dataset in datasets.items()} # subsampled training dataloader for faster logging of training performance dataloaders['Train'] = islice(dataloaders['train'],len(dataloaders['test']))#islice(dataloaders['train'],len(dataloaders['train'])//10) # Initialize optimizer and learning rate schedule opt_constr = functools.partial(Adam, lr=lr) cos = cosLr(num_epochs) lr_sched = lambda e: min(e / (.01 * num_epochs), 1) * cos(e) return MoleculeTrainer(model,dataloaders,opt_constr,lr_sched, task=task,ds_stats=ds_stats,**trainer_config)
'augerino': True, 'aug': True }, 'base': { 'augerino': False, 'aug': False }, 'aug': { 'augerino': False, 'aug': True } } if __name__ == '__main__': Trial = train_trial(makeTrainer) config_spec = copy.deepcopy(makeTrainer.__kwdefaults__) config_spec['trainer_config']['log_suffix'] = 'qm9_augerino' config_spec['trainer_config']['early_stop_metric'] = 'valid_MAE' config_spec.update({'task': ['h**o', 'lumo']}) config_spec['net_config']['group'] = [lieGroups.Trivial(3), lieGroups.T(3)] name = 'qm9_augerino_expt_full' #config_spec.pop('study_name') thestudy = Study(Trial, {}, study_name=name, base_log_dir=config_spec['trainer_config'].get( 'log_dir', None)) for name, net_cfg in expt_settings.items(): the_config = copy.deepcopy(config_spec) the_config['net_config'].update(net_cfg) the_config['name'] = name thestudy.run(num_trials=-1, new_config_spec=the_config, ordered=True) print(thestudy.results_df())