mean, stddev = torch.tensor([-1202.6432]), torch.tensor([12.3304])
print(mean, stddev)
print(mean_forces, stddev_forces)
model = get_model(train_args,
atomref=None,
mean=mean,
stddev=stddev,
train_loader=train_loader,
parallelize=args.parallel)
data = DataBunch(train_loader,
val_loader,
collate_fn=schnetpack2.data.collate_atoms)
learn = Learner(data, model, model_dir=args.modelpath)
#learn.lr_find(start_lr=1e-6,end_lr=1e0, num_it=300, wd=wd)
#plot_recorder(learn.recorder, 'uncertainties_wd_{0}'.format(wd))
#print(hasattr(list(learn.model.children())[-3].weight, "dict"))
#pdb.set_trace()
try:
if args.num_inducing_points:
learn.model.load_state_dict(
torch.load(
os.path.join(
args.modelpath, "bulkVtrain3200" + "_induced_" +
str(sys.argv[3] + ".pth")))['model'])
else:
learn.model.load_state_dict(
print(mean_forces, stddev_forces)
model = get_model(train_args,
atomref=None,
mean=mean,
stddev=stddev,
train_loader=train_loader,
parallelize=args.parallel)
data = DataBunch(train_loader,
val_loader,
collate_fn=schnetpack2.data.collate_atoms)
#pdb.set_trace()
#print(hasattr(list(model.children())[-1].sigma_net[1].out_net[-1].weight, "dict"))
learn = Learner(data, model, model_dir=args.modelpath)
learn.split(split_sigma_net)
#learn.purge()
if args.sdr:
learn.opt_func = Adam_sdr
else:
learn.opt_func = Adam
mean, stddev = torch.tensor([-1202.6432]), torch.tensor([12.3304])
if args.uncertainty_forces:
# learn.loss_func = partial(NLLMSEloss_forces,mean=mean.cuda(), std=stddev.cuda(),kf=100.0,ke=1.0)#best before forces rescaling
learn.loss_func = partial(NLLMSEloss_forces,
mean=mean.cuda(),
if args.split_path is not None:
copyfile(args.split_path, split_path)
train_loader = schnetpack2.custom.data.AtomsLoader(data_train, batch_size=args.batch_size, sampler=RandomSampler(data_train),
num_workers=9*torch.cuda.device_count(), pin_memory=True)
val_loader = schnetpack2.custom.data.AtomsLoader(data_val, batch_size=args.batch_size, num_workers=9*torch.cuda.device_count(), pin_memory=True)
mean, stddev = train_loader.get_statistics('energy', False)
print(mean,stddev)
model = get_model(train_args, atomref=None, mean=torch.FloatTensor([mean]), stddev=torch.FloatTensor([stddev]),
train_loader=train_loader,
parallelize=args.parallel)
data = DataBunch( train_loader, val_loader,collate_fn=schnetpack2.data.collate_atoms)
learn = Learner(data, model, model_dir=args.modelpath)#,callback_fns=ShowGraph)#,callback_fns=SaveAllGrads
learn.opt_func = Adam
learn.loss_func = partial(MSEloss,kf=1,ke=0.1)
learn.metrics=[Emetric,Fmetric]
learn.fit_one_cycle(cyc_len=2, max_lr=3e-3,moms=(0.95, 0.85), div_factor=100.0, pct_start=0.3, wd=1e-7,no_grad_val=False)
learn.to_fp16(loss_scale=100, max_noskip=1000, dynamic=False, clip=None,flat_master=True)
torch.cuda.empty_cache()
#learn.lr_find(start_lr=1e-6,end_lr=1e-2,no_grad_val = False,num_it=63)
#learn.recorder.plot()
copyfile(args.split_path, split_path)
#from sklearn.model_selection import train_test_split
#train,test = train_test_split(df, test_size=0.20, random_state=42,stratify=df['Ebin'].values)
print(args.batch_size)
train_loader = schnetpack2.custom.data.AtomsLoader(data_train, batch_size=args.batch_size, sampler=RandomSampler(data_train),
num_workers=0*torch.cuda.device_count(), pin_memory=True)#9*torch.cuda.device_count()
val_loader = schnetpack2.custom.data.AtomsLoader(data_val, batch_size=args.batch_size, num_workers=9*torch.cuda.device_count(), pin_memory=True)#9*torch.cuda.device_count()
mean, stddev = train_loader.get_statistics('energy', False)
#mean, stddev = torch.tensor([-1.5115]), torch.tensor([1.2643])
#stddev = 29.6958
print(mean,stddev)
model = get_model(train_args, atomref=None, mean=torch.FloatTensor([mean]), stddev=torch.FloatTensor([stddev]),
train_loader=train_loader,
parallelize=args.parallel)
data = DataBunch( train_loader, val_loader,collate_fn=schnetpack2.custom.data.collate_atoms)
learn = Learner(data, model, model_dir=args.modelpath)
learn.purge()
learn.opt_func = Adam
learn.loss_func = partial(MSEloss,kf=1.0,ke=0.1)
learn.metrics=[Emetric,Fmetric]
#learn.load('trainbulkVsurf128epochs')
torch.cuda.empty_cache()
print(sys.argv)
learn.fit_one_cycle(cyc_len=int(sys.argv[3]), max_lr=float(sys.argv[4]),moms=(0.95, 0.85), div_factor=150.0, pct_start=0.05, wd=1e-2,no_grad_val=False)
print(learn.save( sys.argv[1], return_path=True))
print(mean, stddev)
print(mean_forces, stddev_forces)
model = get_model(train_args,
atomref=None,
mean=mean,
stddev=stddev,
train_loader=train_loader,
parallelize=args.parallel)
data = DataBunch(train_loader,
val_loader,
collate_fn=schnetpack2.data.collate_atoms)
#print(hasattr(list(model.children())[-1].sigma_net[1].out_net[-1].weight, "dict"))
learn = Learner(data, model, model_dir=args.modelpath)
learn.split(split_sigma_net)
#learn.purge()
if args.sdr:
learn.opt_func = Adam_sdr
else:
learn.opt_func = Adam
mean, stddev = torch.tensor([-1202.6432]), torch.tensor([12.3304])
#mean, stddev = torch.tensor([-1202.6432,0]), torch.tensor([12.3304,1])
if args.uncertainty_forces:
# learn.loss_func = partial(NLLMSEloss_forces,mean=mean.cuda(), std=stddev.cuda(),kf=100.0,ke=1.0)#best before forces rescaling
learn.loss_func = partial(NLLMSEloss_forces,
mean, stddev = torch.tensor([-1202.6432]), torch.tensor([12.3304])
print(mean, stddev)
print(mean_forces, stddev_forces)
model = get_model(train_args,
atomref=None,
mean=torch.FloatTensor([mean]),
stddev=torch.FloatTensor([stddev]),
train_loader=train_loader,
parallelize=args.parallel)
data = DataBunch(train_loader,
val_loader,
collate_fn=schnetpack2.data.collate_atoms)
learn = Learner(data, model, model_dir=args.modelpath)
#pdb.set_trace()
learn.opt_func = Adam
#learn.loss_func = partial(NLLMSEloss_forces, mean=mean.cuda(), std=stddev.cuda(), mean_forces=mean_forces.cuda(), std_forces=stddev_forces.cuda(), kf=0.001,ke=1)
if args.num_inducing_points:
learn.loss_func = partial(
NLLMSEloss,
mean=mean.cuda(),
std=stddev.cuda(),
kf=0.5,
ke=1.0,
kgamma=0.0001) #partial(MSEloss, kf=0.1,ke=1,kef=0)
else:
learn.loss_func = partial(MSEloss, kf=0.1, ke=1, kef=0)
mean, stddev = train_loader.get_statistics('energy', False)
#mean, stddev = torch.tensor([-1.5115]), torch.tensor([1.2643])
#stddev = 29.6958
print(mean, stddev)
model = get_model(train_args,
atomref=None,
mean=torch.FloatTensor([mean]),
stddev=torch.FloatTensor([stddev]),
train_loader=train_loader,
parallelize=args.parallel)
data = DataBunch(train_loader,
val_loader,
collate_fn=schnetpack2.custom.data.collate_atoms)
learn = Learner(data, model, model_dir=args.modelpath)
learn.purge()
learn.opt_func = Adam
learn.loss_func = partial(MSEloss, kf=1.0, ke=0.1)
learn.metrics = [Emetric, Fmetric]
#learn.load('trainbulkVsurf128epochs')
import matplotlib.pyplot as plt
import numpy as np
def plot_recorder(recorder,
save_name,
skip_start: int = 10,
skip_end: int = 5):
"Plot learning rate and losses, trimmed between `skip_start` and `skip_end`. Optionally plot and return min gradient"