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