def net_builder(self, train_dict):
neval = velocity_verlet3.neval
pwnet_list = []
mbnet_list = []
pw4mb_list = []
output_dim = 2
for n in range(neval):
pwnet_list.append(
mydevice.load(
pw_net(n + 2, output_dim, train_dict["pwnet_nnodes"])))
mbnet_list.append(
mydevice.load(
mb_net(12 * n + 24, output_dim,
train_dict["mbnet_nnodes"])))
for n in range(neval // 2):
pw4mb_list.append(
mydevice.load(pw_net(2, output_dim,
train_dict["pw4mb_nnodes"])))
ngrids = train_dict["ngrids"]
b = train_dict["b"]
pw_obj = pw_ff(pwnet_list)
mb_obj = mb_ff(mbnet_list, pw4mb_list, ngrids, b)
mbpw_obj = mbpw(mb_obj, pw_obj)
# concatenate all network for checkpoint
net_list = pwnet_list + mbnet_list + pw4mb_list
return mbpw_obj, net_list
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)
nnet(72, 2),
nnet(84, 2)
]
pw4mb_list = [nnet(2, 2), nnet(2, 2), nnet(2, 2)]
for pw, mb in zip(pwnet_list, mbnet_list):
pw = mydevice.load(pw)
mb = mydevice.load(mb)
for pw4mb in pw4mb_list:
pw4mb = mydevice.load(pw4mb)
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)