def __init__(self, train_dict, loss_dict):
self.train_dict = train_dict
mode = train_dict["nn_mode"]
if mode == 'ff':
output_dim = 1
else:
print('hf not implemented yet')
quit()
self.mbpw_obj, self.net_list = self.net_builder(train_dict)
self.mlvv = velocity_verlet3(self.mbpw_obj)
self.opt = optim.SGD(self.mbpw_obj.parameters(), lr=train_dict["lr"])
self.sch = optim.lr_scheduler.StepLR(self.opt, train_dict["sch_step"],
train_dict["sch_decay"])
#self.sch = DecayCosineAnnealingWarmRestarts(self.opt,train_dict["sch_step"],train_dict["sch_decay"])
self.opt_tau = optim.SGD(self.mbpw_obj.tau_parameters(),
train_dict["tau_lr"])
lj = lennard_jones2d()
self.loss_obj = loss(lj, loss_dict["eweight"],
loss_dict["polynomial_degree"])
def __init__(self, train_dict, loss_dict):
self.train_dict = train_dict
pwnet_input = 5 # for pw force function
mbnet_input = 25 # for mb force function
mode = train_dict["nn_mode"]
if mode == 'hf':
output_dim = 1
else:
output_dim = 2
pwnet1 = mydevice.load(
pw_net(pwnet_input, output_dim, train_dict["nnodes"]))
pwnet2 = mydevice.load(
pw_net(pwnet_input, output_dim, train_dict["nnodes"]))
mbnet1 = mydevice.load(
mb_net(mbnet_input, output_dim, train_dict["nnodes"]))
mbnet2 = mydevice.load(
mb_net(mbnet_input, output_dim, train_dict["nnodes"]))
self.net_list = [pwnet1, pwnet2, mbnet1, mbnet2]
print('pwnet1', pwnet1)
print('pwnet2', pwnet2)
print('mbnet1', mbnet1)
print('mbnet2', mbnet2)
if mode == 'hf':
pwforce = pw_hf(pwnet1, pwnet2, train_dict["force_clip"])
mbforce = mb_hf(mbnet1, mbnet2, train_dict["ngrids"],
train_dict["b"], train_dict["force_clip"])
else:
pwforce = pw_ff(pwnet1, pwnet2, train_dict["force_clip"])
mbforce = mb_ff(mbnet1, mbnet2, train_dict["ngrids"],
train_dict["b"], train_dict["force_clip"])
self.mbpwff = mbpw(mbforce, pwforce)
param = itertools.chain(pwnet1.parameters(), pwnet2.parameters(),
mbnet1.parameters(), mbnet2.parameters())
self.opt = optim.SGD(param, lr=train_dict["lr"])
#sch = optim.lr_scheduler.StepLR(self.opt,train_dict["sch_step"],train_dict["sch_decay"])
self.sch = DecayCosineAnnealingWarmRestarts(self.opt,
train_dict["sch_step"],
train_dict["sch_decay"])
self.opta = optim.SGD(self.mbpwff.parameters(), train_dict["alpha_lr"])
self.mlvv = velocity_verlet(self.mbpwff)
lj = lennard_jones2d()
self.loss_obj = loss(lj, loss_dict["eweight"],
loss_dict["polynomial_degree"])
self.ckpt = checkpoint(self.net_list, self.mbpwff, self.opt, self.opta,
self.sch)
mb = mb_ff(mbnet_list, pw4mb_list, ngrids, b, force_clip, nsamples,
nparticles)
pw = pw_ff(pwnet_list, force_clip, nsamples, nparticles)
mbpw_obj = mbpw(mb, pw)
vv3 = velocity_verlet3(mbpw_obj)
lr = 1e-4
opt = optim.SGD(mbpw_obj.parameters(), lr)
opt2 = optim.SGD(mbpw_obj.tau_parameters(), lr * 1e-1)
sch = torch.optim.lr_scheduler.StepLR(opt, step_size=1000, gamma=0.9)
energy = potential_function()
poly_deg = 2
loss_obj = loss(energy, 0.0, poly_deg)
nepoch = 100000
for e in range(nepoch):
data_id = np.random.randint(0, nbatches)
q_init = q_init_list[data_id]
p_init = p_init_list[data_id]
l_init = l_init_list[data_id]
q_label = q_label_list[data_id]
p_label = p_label_list[data_id]
q_init = mydevice.load(q_init)
p_init = mydevice.load(p_init)
l_init = mydevice.load(l_init)