},
{
'network': OGN,
'net_cfg': {
'k': 256
},
'lr': 1e-2
},
]
if __name__ == '__main__':
config_spec = copy.deepcopy(makeTrainer.__kwdefaults__)
config_spec.update({
'num_epochs': (lambda cfg: int(np.sqrt(1e7 / cfg['n_train']))),
'n_train':
[10, 25, 50, 100, 400, 1000, 3000, 10000, 30000, 100000 - 4000],
})
config_spec = argupdated_config(config_spec)
name = 'data_scaling_dynamics_final' #config_spec.pop('study_name')
num_repeats = 3 #config_spec.pop('num_repeats')
thestudy = Study(Trial, {},
study_name=name,
base_log_dir=config_spec['trainer_config'].get(
'log_dir', None))
for cfg in best_hypers:
the_config = copy.deepcopy(config_spec)
the_config.update(cfg)
thestudy.run(num_trials=-1 * num_repeats,
new_config_spec=the_config,
ordered=True)
print(thestudy.results_df())
logging.getLogger().setLevel(log_levels[log_level])
# Prep the datasets splits, model, and dataloaders
with FixedNumpySeed(seed), FixedPytorchSeed(seed):
base_dataset = dataset(ndata)
datasets = split_dataset(base_dataset, splits=split)
if net_config['group'] is None: net_config['group'] = base_dataset.symmetry
model = network(base_dataset.rep_in, base_dataset.rep_out, **net_config)
if aug: model = base_dataset.default_aug(model)
model = Standardize(model, datasets['train'].stats)
dataloaders = {
k: LoaderTo(
DataLoader(v,
batch_size=min(bs, len(v)),
shuffle=(k == 'train'),
num_workers=0,
pin_memory=False))
for k, v in datasets.items()
}
dataloaders['Train'] = dataloaders['train']
opt_constr = objax.optimizer.Adam
lr_sched = lambda e: lr #*min(1,e/(num_epochs/10)) # Learning rate warmup
return RegressorPlus(model, dataloaders, opt_constr, lr_sched,
**trainer_config)
if __name__ == "__main__":
cfg = argupdated_config(makeTrainer.__kwdefaults__,
namespace=(emlp.groups, emlp.datasets, emlp.nn))
trainer = makeTrainer(**cfg)
trainer.train(cfg['num_epochs'])
import torch
import lie_conv.lieGroups as lieGroups
import functools
import copy
import pandas as pd
from augerino.models.qm9_models import makeTrainer
import copy
from oil.tuning.study import train_trial,Study
from oil.tuning.args import argupdated_config
import lie_conv.moleculeTrainer as moleculeTrainer
# Example run single with argument parsing
if __name__=='__main__':
Trial = train_trial(makeTrainer)
defaults = copy.deepcopy(makeTrainer.__kwdefaults__)
defaults['trainer_config']['early_stop_metric']='valid_MAE'
defaults['save']=False
cfg,outcome = Trial(argupdated_config(defaults,namespace=(moleculeTrainer,lieGroups)))
print(outcome)
}
# subsampled training dataloader for faster logging of training performance
dataloaders['Train'] = 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)
Trial = train_trial(makeTrainer)
if __name__ == '__main__':
defaults = copy.deepcopy(makeTrainer.__kwdefaults__)
defaults['trainer_config']['early_stop_metric'] = 'Train_MAE'
print(
Trial(
argupdated_config(defaults,
namespace=(moleculeTrainer, lieGroups))))
# thestudy = Study(simpleTrial,argupdated_config(config_spec,namespace=__init__),
# study_name="point2d",base_log_dir=log_dir)
# thestudy.run(ordered=False)
model = network(num_classes=datasets['train'].num_classes,
dim_in=datasets['train'].dim,
**net_config).to(device)
dataloaders = {
k: LoaderTo(
DataLoader(v,
batch_size=min(bs, len(datasets[k])),
shuffle=(k == 'train'),
num_workers=0,
pin_memory=False), device)
for k, v in datasets.items()
}
dataloaders['Train'] = dataloaders['train']
opt_constr = partial(optim, lr=lr, **opt_config)
lr_sched = cosLr(num_epochs) #lambda e:1#
return trainer(model, dataloaders, opt_constr, lr_sched, **trainer_config)
tabularTrial = train_trial(makeTrainer)
if __name__ == '__main__':
defaults = copy.deepcopy(makeTrainer.__kwdefaults__)
cfg = argupdated_config(defaults,
namespace=(tabular_datasets, flows, archs,
trainers))
cfg.pop('local_rank')
trainer = makeTrainer(**cfg)
#tabularTrial()
trainer.train(cfg['num_epochs'])
return IntegratedDynamicsTrainer(model,
dataloaders,
opt_constr,
lr_sched,
log_args={
'timeFrac': 1 / 4,
'minPeriod': 0.0
},
**trainer_config)
# best_hypers = {
# LieResNetT2: {'net_cfg':{'k':384, 'num_layers':4},'lr':1e-3},
# VOGN: {'net_cfg':{'k':512},'lr':3e-3},
# HOGN: {'net_cfg':{'k':256},'lr':1e-2},
# OGN: {'net_cfg':{'k':256},'lr':1e-2},
# FCHamNet: {'net_cfg':{'k':256,'num_layers':4},'lr':1e-2},
# RawDynamicsNet: {'net_cfg':{'k':256},'lr':3e-3},
# }
Trial = DynamicsTrial(makeTrainer)
if __name__ == '__main__':
defaults = copy.deepcopy(makeTrainer.__kwdefaults__)
#defaults.update({'network':HLieResNet,'net_cfg':{'group':Tx(2),'k':384,'num_layers':4},'lr':1e-3})
#defaults['early_stop_metric']='val_MSE'
#Trial(defaults)
print(
Trial(
argupdated_config(defaults,
namespace=(dynamicsTrainer, lieGroups, datasets,
graphnets))))
# Prep the datasets splits, model, and dataloaders
with FixedNumpySeed(seed), FixedPytorchSeed(seed):
datasets = {
'train': InvertedCube(train=True),
'test': InvertedCube(train=False)
}
model = Standardize(
network(datasets['train'].rep_in, datasets['train'].rep_out,
**net_config), datasets['train'].stats)
dataloaders = {
k: LoaderTo(
DataLoader(v,
batch_size=min(bs, len(v)),
shuffle=(k == 'train'),
num_workers=0,
pin_memory=False))
for k, v in datasets.items()
}
dataloaders['Train'] = dataloaders['train']
opt_constr = objax.optimizer.Adam
lr_sched = lambda e: lr * cosLr(num_epochs)(e)
return ClassifierPlus(model, dataloaders, opt_constr, lr_sched,
**trainer_config)
if __name__ == "__main__":
Trial = train_trial(makeTrainer)
Trial(
argupdated_config(makeTrainer.__kwdefaults__,
namespace=(emlp.solver.groups, emlp.models.datasets,
emlp.models.mlp)))
opt_constr,
lr_sched,
log_args={
"timeFrac": 1 / 4,
"minPeriod": 0.0
},
**trainer_config)
#Trial = train_trial(makeTrainer)
if __name__ == "__main__":
with FixedNumpySeed(0):
defaults = copy.deepcopy(makeTrainer.__kwdefaults__)
defaults["save"] = False
namespace = (datasets, systems, models)
cfg = argupdated_config(defaults, namespace=namespace)
cfg.pop('local_rank')
save = cfg.pop('save')
trainer = makeTrainer(**cfg)
trainer.train(cfg['num_epochs'])
if save: print(f"saved at: {trainer.save_checkpoint()}")
rollouts = trainer.test_rollouts(
angular_to_euclidean=not issubclass(cfg['network'], (CH, CL)))
print(
f"rollout error GeoMean {rollouts[0][:,1:].log().mean().exp():.3E}"
)
fname = f"rollout_errs_{cfg['network']}_{cfg['body']}.np"
with open(fname, 'wb') as f:
pickle.dump(rollouts, f)
#defaults["trainer_config"]["early_stop_metric"] = "val_MSE"
#print(Trial()))
import lie_conv.moleculeTrainer as moleculeTrainer
import lie_conv.lieGroups as lieGroups
from lie_conv.lieGroups import T,Trivial,SE3,SO3
import lie_conv.lieConv as lieConv
from lie_conv.lieConv import ImgLieResnet
from lie_conv.datasets import MnistRotDataset
from examples.train_molec import makeTrainer,Trial
from oil.tuning.study import Study
def trial_name(cfg):
ncfg = cfg['net_config']
return f"molec_f{ncfg['fill']}_n{ncfg['nbhd']}_{ncfg['group']}_{cfg['lr']}"
def bigG(cfg):
return isinstance(cfg['net_config']['group'],(SE3,SO3))
if __name__ == '__main__':
config_spec = copy.deepcopy(makeTrainer.__kwdefaults__)
config_spec.update({
'num_epochs':500,
'net_config':{'fill':lambda cfg: (1.,1/2)[bigG(cfg)],'nbhd':lambda cfg: (100,25)[bigG(cfg)],
'group':T(3),'liftsamples':lambda cfg: (1,4)[bigG(cfg)]},'recenter':lambda cfg: bigG(cfg),
'lr':3e-3,'bs':lambda cfg: (100,75)[bigG(cfg)],'task':['alpha','gap','h**o','lumo','mu','Cv','G','H','r2','U','U0','zpve'],
'trainer_config':{'log_dir':'molec_all_tasks4','log_suffix':lambda cfg:trial_name(cfg)},
})
config_spec = argupdated_config(config_spec,namespace=(moleculeTrainer,lieGroups))
thestudy = Study(Trial,config_spec,study_name='molec_all_tasks4')
thestudy.run(num_trials=-1,ordered=True)
print(thestudy.results_df())
# Run MLP baseline on datasets
config_spec.update({
'dataset':
ParticleInteraction, #[O5Synthetic,Inertia,ParticleInteraction],
'network':
MLP,
'aug': [False, True],
'num_epochs':
(lambda cfg: min(int(30 * 30000 / cfg['split']['train']), 1000)),
'split': {
'train': [30, 100, 300, 1000, 3000, 10000, 30000],
'test': 5000,
'val': 1000
},
})
config_spec = argupdated_config(config_spec, namespace=datasets.regression)
name = f"{name}_{config_spec['dataset']}"
thestudy = Study(Trial, {},
study_name=name,
base_log_dir=config_spec['trainer_config'].get(
'log_dir', None))
thestudy.run(num_trials=-3, new_config_spec=config_spec, ordered=True)
# Now run the EMLP (with appropriate group) on the datasets
config_spec['network'] = EMLP
config_spec['aug'] = False
groups = {
O5Synthetic: [SO(5), O(5)],
Inertia: [SO(3), O(3)],
ParticleInteraction: [SO13p(), SO13(), O13()]
}
if aug:
model = torch.nn.Sequential(datasets['train'].default_aug_layers(),
model)
model, bs = try_multigpu_parallelize(model, bs)
dataloaders = {
k: LoaderTo(
DataLoader(v,
batch_size=bs,
shuffle=(k == 'train'),
num_workers=0,
pin_memory=False), device)
for k, v in datasets.items()
}
dataloaders['Train'] = islice(dataloaders['train'],
1 + len(dataloaders['train']) // 10)
if small_test:
dataloaders['test'] = islice(dataloaders['test'],
1 + len(dataloaders['train']) // 10)
# Add some extra defaults if SGD is chosen
opt_constr = partial(optim, lr=lr, **opt_config)
lr_sched = cosLr(num_epochs)
return trainer(model, dataloaders, opt_constr, lr_sched, **trainer_config)
if __name__ == "__main__":
Trial = train_trial(makeTrainer)
defaults = copy.deepcopy(makeTrainer.__kwdefaults__)
defaults['save'] = False
Trial(argupdated_config(defaults, namespace=(lieConv, lieGroups)))