def train (args) :
# load json database
with open (args.INPUT, 'r') as fp:
jdata = json.load (fp)
if not 'model' in jdata.keys():
jdata = convert_input_v0_v1(jdata,
warning = True,
dump = 'input_v1_compat.json')
# run options
with_distrib = False
if 'with_distrib' in jdata:
with_distrib = jdata['with_distrib']
run_opt = RunOptions(args, with_distrib)
run_opt.print_welcome()
run_opt.print_citation()
run_opt.print_summary()
if run_opt.is_distrib :
# distributed training
if run_opt.my_job_name == "ps":
queue = create_done_queue(run_opt.cluster_spec, run_opt.my_task_index)
wait_done_queue(run_opt.cluster_spec, run_opt.server, queue, run_opt.my_task_index)
#server.join()
elif run_opt.my_job_name == "worker":
done_ops = connect_done_queue(run_opt.cluster_spec, run_opt.my_task_index)
_do_work(jdata, run_opt)
fill_done_queue(run_opt.cluster_spec, run_opt.server, done_ops, run_opt.my_task_index)
else :
raise RuntimeError("unknown job name")
else :
# serial training
_do_work(jdata, run_opt)
def __init__(self):
# setup
with open(os.path.join(os.path.dirname(__file__),
'water_se_a.json')) as fp:
jdata = json.load(fp)
self.run_opt = RunOptions(None)
self.run_opt.verbose = False
self.model = NNPTrainer(jdata, run_opt=self.run_opt)
rcut = self.model.model.get_rcut()
type_map = self.model.model.get_type_map()
systems = [os.path.join(os.path.dirname(__file__), 'data')]
set_pfx = jdata['training']['set_prefix']
seed = jdata['training']['seed']
np.random.seed(seed)
batch_size = jdata['training']['batch_size']
test_size = jdata['training']['numb_test']
self.data = DeepmdDataSystem(systems,
batch_size,
test_size,
rcut,
set_prefix=set_pfx,
run_opt=self.run_opt,
type_map=type_map)
self.data.add_dict(data_requirement)
self.model.build(self.data)
self.model._init_sess_serial()
cur_batch = self.model.sess.run(self.model.global_step)
self.cur_batch = cur_batch
def compute_efv(jfile):
fp = open(jfile, 'r')
jdata = json.load(fp)
run_opt = RunOptions(None)
systems = j_must_have(jdata, 'systems')
set_pfx = j_must_have(jdata, 'set_prefix')
batch_size = j_must_have(jdata, 'batch_size')
test_size = j_must_have(jdata, 'numb_test')
batch_size = 1
test_size = 1
rcut = j_must_have(jdata, 'rcut')
data = DataSystem(systems, set_pfx, batch_size, test_size, rcut, run_opt)
tot_numb_batches = sum(data.get_nbatches())
lr = LearingRate(jdata, tot_numb_batches)
model = NNPModel(jdata, run_opt=run_opt)
model.build(data, lr)
test_data = data.get_test()
feed_dict_test = {
model.t_prop_c:
test_data["prop_c"],
model.t_energy:
test_data["energy"][:model.numb_test],
model.t_force:
np.reshape(test_data["force"][:model.numb_test, :], [-1]),
model.t_virial:
np.reshape(test_data["virial"][:model.numb_test, :], [-1]),
model.t_atom_ener:
np.reshape(test_data["atom_ener"][:model.numb_test, :], [-1]),
model.t_atom_pref:
np.reshape(test_data["atom_pref"][:model.numb_test, :], [-1]),
model.t_coord:
np.reshape(test_data["coord"][:model.numb_test, :], [-1]),
model.t_box:
test_data["box"][:model.numb_test, :],
model.t_type:
np.reshape(test_data["type"][:model.numb_test, :], [-1]),
model.t_natoms:
test_data["natoms_vec"],
model.t_mesh:
test_data["default_mesh"],
model.t_fparam:
np.reshape(test_data["fparam"][:model.numb_test, :], [-1]),
model.is_training:
False
}
sess = tf.Session()
sess.run(tf.global_variables_initializer())
[e, f, v] = sess.run([model.energy, model.force, model.virial],
feed_dict=feed_dict_test)
return e, f, v
def _setUp(self):
args = Args()
run_opt = RunOptions(args, False)
with open (args.INPUT, 'r') as fp:
jdata = json.load (fp)
# init model
model = NNPTrainer (jdata, run_opt = run_opt)
rcut = model.model.get_rcut()
# init data system
systems = j_must_have(jdata['training'], 'systems')
set_pfx = j_must_have(jdata['training'], 'set_prefix')
batch_size = j_must_have(jdata['training'], 'batch_size')
test_size = j_must_have(jdata['training'], 'numb_test')
data = DeepmdDataSystem(systems,
batch_size,
test_size,
rcut,
set_prefix=set_pfx)
data.add_dict(data_requirement)
# clear the default graph
tf.reset_default_graph()
# build the model with stats from the first system
model.build (data)
# freeze the graph
with tf.Session() as sess:
init_op = tf.global_variables_initializer()
sess.run(init_op)
graph = tf.get_default_graph()
input_graph_def = graph.as_graph_def()
nodes = "o_dipole,o_rmat,o_rmat_deriv,o_nlist,o_rij,descrpt_attr/rcut,descrpt_attr/ntypes,descrpt_attr/sel,descrpt_attr/ndescrpt,model_attr/tmap,model_attr/sel_type,model_attr/model_type"
output_graph_def = tf.graph_util.convert_variables_to_constants(
sess,
input_graph_def,
nodes.split(",")
)
output_graph = os.path.join(modifier_datapath, 'dipole.pb')
with tf.gfile.GFile(output_graph, "wb") as f:
f.write(output_graph_def.SerializeToString())
def __init__(self,
sess,
jdata,
run_opt = RunOptions()):
self.run_opt = run_opt
self.sess = sess
# descrpt config
self.use_smooth = False
if j_have (jdata, "use_smooth") :
self.use_smooth = jdata["use_smooth"]
self.sel_a = j_must_have (jdata, 'sel_a')
self.sel_r = [ 0 for ii in range(len(self.sel_a)) ]
if not self.use_smooth :
self.sel_r = j_must_have (jdata, 'sel_r')
else :
if j_have (jdata, 'sel_r') :
self.warning ('ignoring key sel_r in the json database and set sel_r to %s' %
str(self.sel_r))
self.rcut_a = -1
self.rcut_r = j_must_have (jdata, 'rcut')
if j_have(jdata, 'rcut_smth') :
self.rcut_r_smth = jdata['rcut_smth']
else :
self.rcut_r_smth = self.rcut_r
# axis
self.axis_rule = []
if j_have (jdata, 'axis_rule') :
self.axis_rule = jdata['axis_rule']
# filter of smooth version
if self.use_smooth :
self.filter_neuron = j_must_have (jdata, 'filter_neuron')
self.n_axis_neuron = j_must_have (jdata, 'n_axis_neuron')
self.filter_resnet_dt = False
if j_have(jdata, 'filter_resnet_dt') :
self.filter_resnet_dt = jdata['filter_resnet_dt']
# numb of neighbors and numb of descrptors
self.nnei_a = np.cumsum(self.sel_a)[-1]
self.nnei_r = np.cumsum(self.sel_r)[-1]
self.nnei = self.nnei_a + self.nnei_r
self.ndescrpt_a = self.nnei_a * 4
self.ndescrpt_r = self.nnei_r * 1
self.ndescrpt = self.ndescrpt_a + self.ndescrpt_r
# network size
self.n_neuron = j_must_have (jdata, 'n_neuron')
self.resnet_dt = True
if j_have(jdata, 'resnet_dt') :
self.resnet_dt = jdata['resnet_dt']
self.numb_test = j_must_have (jdata, 'numb_test')
self.useBN = False
self.start_pref_e = j_must_have (jdata, 'start_pref_e')
self.limit_pref_e = j_must_have (jdata, 'limit_pref_e')
self.start_pref_f = j_must_have (jdata, 'start_pref_f')
self.limit_pref_f = j_must_have (jdata, 'limit_pref_f')
self.start_pref_v = j_must_have (jdata, 'start_pref_v')
self.limit_pref_v = j_must_have (jdata, 'limit_pref_v')
self.has_e = (self.start_pref_e != 0 or self.limit_pref_e != 0)
self.has_f = (self.start_pref_f != 0 or self.limit_pref_f != 0)
self.has_v = (self.start_pref_v != 0 or self.limit_pref_v != 0)
self.disp_file = "lcurve.out"
if j_have (jdata, "disp_file") : self.disp_file = jdata["disp_file"]
self.disp_freq = j_must_have (jdata, 'disp_freq')
self.save_freq = j_must_have (jdata, 'save_freq')
self.save_ckpt = j_must_have (jdata, 'save_ckpt')
self.seed = None
if j_have (jdata, 'seed') :
self.seed = jdata['seed']
self.display_in_training = j_must_have (jdata, 'disp_training')
self.timing_in_training = j_must_have (jdata, 'time_training')
self.profiling = False
if j_have (jdata, 'profiling') :
self.profiling = jdata['profiling']
if self.profiling :
self.profiling_file = j_must_have (jdata, 'profiling_file')
self.null_mesh = tf.constant ([-1])
self.verbose = True
def _main () :
default_num_inter_threads = 0
parser = argparse.ArgumentParser(
description="*** Train a model. ***")
parser.add_argument('INPUT',
help='the input json database ')
parser.add_argument('-t','--inter-threads', type = int, default = default_num_inter_threads,
help=
'With default value %d. ' % default_num_inter_threads +
'Setting the "inter_op_parallelism_threads" key for the tensorflow, ' +
'the "intra_op_parallelism_threads" will be set by the env variable OMP_NUM_THREADS')
parser.add_argument('--init-model', type = str,
help=
'Initialize the model by the provided checkpoint.')
parser.add_argument('--restart', type = str,
help=
'Restart the training from the provided checkpoint.')
args = parser.parse_args()
# load json database
fp = open (args.INPUT, 'r')
jdata = json.load (fp)
# init params and run options
systems = j_must_have(jdata, 'systems')
set_pfx = j_must_have(jdata, 'set_prefix')
numb_sys = len(systems)
seed = None
if 'seed' in jdata.keys() : seed = jdata['seed']
batch_size = j_must_have(jdata, 'batch_size')
test_size = j_must_have(jdata, 'numb_test')
stop_batch = j_must_have(jdata, 'stop_batch')
rcut = j_must_have (jdata, 'rcut')
print ("#")
print ("# find %d system(s): " % numb_sys)
data = DataSystem(systems, set_pfx, batch_size, test_size, rcut)
print ("#")
tot_numb_batches = sum(data.get_nbatches())
lr = LearingRate (jdata, tot_numb_batches)
final_lr = lr.value (stop_batch)
run_opt = RunOptions(args)
print("# run with intra_op_parallelism_threads = %d, inter_op_parallelism_threads = %d " %
(run_opt.num_intra_threads, run_opt.num_inter_threads))
# start tf
tf.reset_default_graph()
with tf.Session(
config=tf.ConfigProto(intra_op_parallelism_threads=run_opt.num_intra_threads,
inter_op_parallelism_threads=run_opt.num_inter_threads
)) as sess:
# init the model
model = NNPModel (sess, jdata, run_opt = run_opt)
# build the model with stats from the first system
model.build (data, lr)
# train the model with the provided systems in a cyclic way
start_time = time.time()
cur_batch = model.get_global_step()
print ("# start training, start lr is %e, final lr will be %e" % (lr.value(cur_batch), final_lr) )
model.print_head()
model.train (data, stop_batch)
print ("# finished training")
end_time = time.time()
print ("# running time: %.3f s" % (end_time-start_time))
def test_model(self):
jfile = 'water.json'
with open(jfile) as fp:
jdata = json.load(fp)
run_opt = RunOptions(None)
systems = j_must_have(jdata, 'systems')
set_pfx = j_must_have(jdata, 'set_prefix')
batch_size = j_must_have(jdata, 'batch_size')
test_size = j_must_have(jdata, 'numb_test')
batch_size = 1
test_size = 1
stop_batch = j_must_have(jdata, 'stop_batch')
rcut = j_must_have(jdata['model']['descriptor'], 'rcut')
data = DataSystem(systems,
set_pfx,
batch_size,
test_size,
rcut,
run_opt=None)
test_data = data.get_test()
numb_test = 1
descrpt = DescrptLocFrame(jdata['model']['descriptor'])
fitting = EnerFitting(jdata['model']['fitting_net'], descrpt)
model = Model(jdata['model'], descrpt, fitting)
# model._compute_dstats([test_data['coord']], [test_data['box']], [test_data['type']], [test_data['natoms_vec']], [test_data['default_mesh']])
input_data = {
'coord': [test_data['coord']],
'box': [test_data['box']],
'type': [test_data['type']],
'natoms_vec': [test_data['natoms_vec']],
'default_mesh': [test_data['default_mesh']]
}
model._compute_input_stat(input_data)
model.fitting.bias_atom_e = data.compute_energy_shift()
t_prop_c = tf.placeholder(tf.float32, [5], name='t_prop_c')
t_energy = tf.placeholder(global_ener_float_precision, [None],
name='t_energy')
t_force = tf.placeholder(global_tf_float_precision, [None],
name='t_force')
t_virial = tf.placeholder(global_tf_float_precision, [None],
name='t_virial')
t_atom_ener = tf.placeholder(global_tf_float_precision, [None],
name='t_atom_ener')
t_coord = tf.placeholder(global_tf_float_precision, [None],
name='i_coord')
t_type = tf.placeholder(tf.int32, [None], name='i_type')
t_natoms = tf.placeholder(tf.int32, [model.ntypes + 2],
name='i_natoms')
t_box = tf.placeholder(global_tf_float_precision, [None, 9],
name='i_box')
t_mesh = tf.placeholder(tf.int32, [None], name='i_mesh')
is_training = tf.placeholder(tf.bool)
t_fparam = None
model_pred \
= model.build (t_coord,
t_type,
t_natoms,
t_box,
t_mesh,
t_fparam,
suffix = "loc_frame",
reuse = False)
energy = model_pred['energy']
force = model_pred['force']
virial = model_pred['virial']
atom_ener = model_pred['atom_ener']
feed_dict_test = {
t_prop_c: test_data['prop_c'],
t_energy: test_data['energy'][:numb_test],
t_force: np.reshape(test_data['force'][:numb_test, :], [-1]),
t_virial: np.reshape(test_data['virial'][:numb_test, :], [-1]),
t_atom_ener: np.reshape(test_data['atom_ener'][:numb_test, :],
[-1]),
t_coord: np.reshape(test_data['coord'][:numb_test, :], [-1]),
t_box: test_data['box'][:numb_test, :],
t_type: np.reshape(test_data['type'][:numb_test, :], [-1]),
t_natoms: test_data['natoms_vec'],
t_mesh: test_data['default_mesh'],
is_training: False
}
sess = tf.Session()
sess.run(tf.global_variables_initializer())
[e, f, v] = sess.run([energy, force, virial], feed_dict=feed_dict_test)
e = e.reshape([-1])
f = f.reshape([-1])
v = v.reshape([-1])
refe = [1.165945032784766511e+01]
reff = [
2.356319331246305437e-01, 1.772322096063349284e-01,
1.455439548950788684e-02, 1.968599426000810226e-01,
2.648214484898352983e-01, 7.595232354012236564e-02,
-2.121321856338151401e-01, -2.463886119018566037e-03,
-2.075636300914874069e-02, -9.360310077571798101e-03,
-1.751965198776750943e-01, -2.046405309983102827e-02,
-1.990194093283037535e-01, -1.828347741191920298e-02,
-6.916374506995154325e-02, -1.197997068502068031e-02,
-2.461097746875573200e-01, 1.987744214930105627e-02
]
refv = [
-4.998509978510510265e-01, -1.966169437179327711e-02,
1.136130543869883977e-02, -1.966169437179334650e-02,
-4.575353297894450555e-01, -2.668666556859019493e-03,
1.136130543869887100e-02, -2.668666556859039876e-03,
2.455466940358383508e-03
]
refe = np.reshape(refe, [-1])
reff = np.reshape(reff, [-1])
refv = np.reshape(refv, [-1])
places = 10
for ii in range(e.size):
self.assertAlmostEqual(e[ii], refe[ii], places=places)
for ii in range(f.size):
self.assertAlmostEqual(f[ii], reff[ii], places=places)
for ii in range(v.size):
self.assertAlmostEqual(v[ii], refv[ii], places=places)
def test_model(self):
jfile = 'wfc.json'
with open(jfile) as fp:
jdata = json.load(fp)
run_opt = RunOptions(None)
systems = j_must_have(jdata, 'systems')
set_pfx = j_must_have(jdata, 'set_prefix')
batch_size = j_must_have(jdata, 'batch_size')
test_size = j_must_have(jdata, 'numb_test')
batch_size = 1
test_size = 1
stop_batch = j_must_have(jdata, 'stop_batch')
rcut = j_must_have(jdata['model']['descriptor'], 'rcut')
data = DataSystem(systems,
set_pfx,
batch_size,
test_size,
rcut,
run_opt=None)
test_data = data.get_test()
numb_test = 1
descrpt = DescrptLocFrame(jdata['model']['descriptor'])
fitting = WFCFitting(jdata['model']['fitting_net'], descrpt)
model = WFCModel(jdata['model'], descrpt, fitting)
input_data = {
'coord': [test_data['coord']],
'box': [test_data['box']],
'type': [test_data['type']],
'natoms_vec': [test_data['natoms_vec']],
'default_mesh': [test_data['default_mesh']],
'fparam': [test_data['fparam']],
}
model._compute_dstats(input_data)
t_prop_c = tf.placeholder(tf.float32, [5], name='t_prop_c')
t_energy = tf.placeholder(global_ener_float_precision, [None],
name='t_energy')
t_force = tf.placeholder(global_tf_float_precision, [None],
name='t_force')
t_virial = tf.placeholder(global_tf_float_precision, [None],
name='t_virial')
t_atom_ener = tf.placeholder(global_tf_float_precision, [None],
name='t_atom_ener')
t_coord = tf.placeholder(global_tf_float_precision, [None],
name='i_coord')
t_type = tf.placeholder(tf.int32, [None], name='i_type')
t_natoms = tf.placeholder(tf.int32, [model.ntypes + 2],
name='i_natoms')
t_box = tf.placeholder(global_tf_float_precision, [None, 9],
name='i_box')
t_mesh = tf.placeholder(tf.int32, [None], name='i_mesh')
is_training = tf.placeholder(tf.bool)
t_fparam = None
model_pred \
= model.build (t_coord,
t_type,
t_natoms,
t_box,
t_mesh,
t_fparam,
suffix = "wfc",
reuse = False)
wfc = model_pred['wfc']
feed_dict_test = {
t_prop_c: test_data['prop_c'],
t_coord: np.reshape(test_data['coord'][:numb_test, :], [-1]),
t_box: test_data['box'][:numb_test, :],
t_type: np.reshape(test_data['type'][:numb_test, :], [-1]),
t_natoms: test_data['natoms_vec'],
t_mesh: test_data['default_mesh'],
is_training: False
}
sess = tf.Session()
sess.run(tf.global_variables_initializer())
[p] = sess.run([wfc], feed_dict=feed_dict_test)
p = p.reshape([-1])
refp = [
-9.105016838228578990e-01, 7.196284362034099935e-01,
-9.548516928185298014e-02, 2.764615027095288724e+00,
2.661319598995644520e-01, 7.579512949131941846e-02,
-2.107409067376114997e+00, -1.299080016614967414e-01,
-5.962778584850070285e-01, 2.913899917663253514e-01,
-1.226917174638697094e+00, 1.829523069930876655e+00,
1.015704024959750873e+00, -1.792333611099589386e-01,
5.032898080485321834e-01, 1.808561721292949453e-01,
2.468863482075112081e+00, -2.566442546384765100e-01,
-1.467453783795173994e-01, -1.822963931552128658e+00,
5.843600156865462747e-01, -1.493875280832117403e+00,
1.693322352814763398e-01, -1.877325443995481624e+00
]
places = 6
for ii in range(p.size):
self.assertAlmostEqual(p[ii], refp[ii], places=places)
def test_model(self):
jfile = 'water_se_a.json'
with open(jfile) as fp:
jdata = json.load(fp)
run_opt = RunOptions(None)
systems = j_must_have(jdata, 'systems')
set_pfx = j_must_have(jdata, 'set_prefix')
batch_size = j_must_have(jdata, 'batch_size')
test_size = j_must_have(jdata, 'numb_test')
batch_size = 1
test_size = 1
stop_batch = j_must_have(jdata, 'stop_batch')
rcut = j_must_have(jdata['model']['descriptor'], 'rcut')
data = DataSystem(systems,
set_pfx,
batch_size,
test_size,
rcut,
run_opt=None)
test_data = data.get_test()
numb_test = 1
descrpt = DescrptSeA(jdata['model']['descriptor'])
fitting = EnerFitting(jdata['model']['fitting_net'], descrpt)
model = Model(jdata['model'], descrpt, fitting)
# model._compute_dstats([test_data['coord']], [test_data['box']], [test_data['type']], [test_data['natoms_vec']], [test_data['default_mesh']])
input_data = {
'coord': [test_data['coord']],
'box': [test_data['box']],
'type': [test_data['type']],
'natoms_vec': [test_data['natoms_vec']],
'default_mesh': [test_data['default_mesh']]
}
model._compute_dstats(input_data)
model.bias_atom_e = data.compute_energy_shift()
t_prop_c = tf.placeholder(tf.float32, [5], name='t_prop_c')
t_energy = tf.placeholder(global_ener_float_precision, [None],
name='t_energy')
t_force = tf.placeholder(global_tf_float_precision, [None],
name='t_force')
t_virial = tf.placeholder(global_tf_float_precision, [None],
name='t_virial')
t_atom_ener = tf.placeholder(global_tf_float_precision, [None],
name='t_atom_ener')
t_coord = tf.placeholder(global_tf_float_precision, [None],
name='i_coord')
t_type = tf.placeholder(tf.int32, [None], name='i_type')
t_natoms = tf.placeholder(tf.int32, [model.ntypes + 2],
name='i_natoms')
t_box = tf.placeholder(global_tf_float_precision, [None, 9],
name='i_box')
t_mesh = tf.placeholder(tf.int32, [None], name='i_mesh')
is_training = tf.placeholder(tf.bool)
t_fparam = None
model_pred \
= model.build (t_coord,
t_type,
t_natoms,
t_box,
t_mesh,
t_fparam,
suffix = "se_a",
reuse = False)
energy = model_pred['energy']
force = model_pred['force']
virial = model_pred['virial']
atom_ener = model_pred['atom_ener']
feed_dict_test = {
t_prop_c: test_data['prop_c'],
t_energy: test_data['energy'][:numb_test],
t_force: np.reshape(test_data['force'][:numb_test, :], [-1]),
t_virial: np.reshape(test_data['virial'][:numb_test, :], [-1]),
t_atom_ener: np.reshape(test_data['atom_ener'][:numb_test, :],
[-1]),
t_coord: np.reshape(test_data['coord'][:numb_test, :], [-1]),
t_box: test_data['box'][:numb_test, :],
t_type: np.reshape(test_data['type'][:numb_test, :], [-1]),
t_natoms: test_data['natoms_vec'],
t_mesh: test_data['default_mesh'],
is_training: False
}
sess = tf.Session()
sess.run(tf.global_variables_initializer())
[e, f, v] = sess.run([energy, force, virial], feed_dict=feed_dict_test)
e = e.reshape([-1])
f = f.reshape([-1])
v = v.reshape([-1])
refe = [6.135449167779321300e+01]
reff = [
7.799691562262310585e-02, 9.423098804815030483e-02,
3.790560997388224204e-03, 1.432522403799846578e-01,
1.148392791403983204e-01, -1.321871172563671148e-02,
-7.318966526325138000e-02, 6.516069212737778116e-02,
5.406418483320515412e-04, 5.870713761026503247e-02,
-1.605402669549013672e-01, -5.089516979826595386e-03,
-2.554593467731766654e-01, 3.092063507347833987e-02,
1.510355029451411479e-02, 4.869271842355533952e-02,
-1.446113274345035005e-01, -1.126524434771078789e-03
]
refv = [
-6.076776685178300053e-01, 1.103174323630009418e-01,
1.984250991380156690e-02, 1.103174323630009557e-01,
-3.319759402259439551e-01, -6.007404107650986258e-03,
1.984250991380157036e-02, -6.007404107650981921e-03,
-1.200076017439753642e-03
]
refe = np.reshape(refe, [-1])
reff = np.reshape(reff, [-1])
refv = np.reshape(refv, [-1])
places = 10
for ii in range(e.size):
self.assertAlmostEqual(e[ii], refe[ii], places=places)
for ii in range(f.size):
self.assertAlmostEqual(f[ii], reff[ii], places=places)
for ii in range(v.size):
self.assertAlmostEqual(v[ii], refv[ii], places=places)
def test_model(self):
jfile = 'polar_se_a.json'
with open(jfile) as fp:
jdata = json.load(fp)
run_opt = RunOptions(None)
systems = j_must_have(jdata, 'systems')
set_pfx = j_must_have(jdata, 'set_prefix')
batch_size = j_must_have(jdata, 'batch_size')
test_size = j_must_have(jdata, 'numb_test')
batch_size = 1
test_size = 1
stop_batch = j_must_have(jdata, 'stop_batch')
rcut = j_must_have(jdata['model']['descriptor'], 'rcut')
data = DataSystem(systems,
set_pfx,
batch_size,
test_size,
rcut,
run_opt=None)
test_data = data.get_test()
numb_test = 1
descrpt = DescrptSeA(jdata['model']['descriptor'])
fitting = PolarFittingSeA(jdata['model']['fitting_net'], descrpt)
model = PolarModel(jdata['model'], descrpt, fitting)
model._compute_dstats([test_data['coord']], [test_data['box']],
[test_data['type']], [test_data['natoms_vec']],
[test_data['default_mesh']])
t_prop_c = tf.placeholder(tf.float32, [5], name='t_prop_c')
t_energy = tf.placeholder(global_ener_float_precision, [None],
name='t_energy')
t_force = tf.placeholder(global_tf_float_precision, [None],
name='t_force')
t_virial = tf.placeholder(global_tf_float_precision, [None],
name='t_virial')
t_atom_ener = tf.placeholder(global_tf_float_precision, [None],
name='t_atom_ener')
t_coord = tf.placeholder(global_tf_float_precision, [None],
name='i_coord')
t_type = tf.placeholder(tf.int32, [None], name='i_type')
t_natoms = tf.placeholder(tf.int32, [model.ntypes + 2],
name='i_natoms')
t_box = tf.placeholder(global_tf_float_precision, [None, 9],
name='i_box')
t_mesh = tf.placeholder(tf.int32, [None], name='i_mesh')
is_training = tf.placeholder(tf.bool)
t_fparam = None
model_pred \
= model.build (t_coord,
t_type,
t_natoms,
t_box,
t_mesh,
t_fparam,
suffix = "polar_se_a",
reuse = False)
polar = model_pred['polar']
feed_dict_test = {
t_prop_c: test_data['prop_c'],
t_coord: np.reshape(test_data['coord'][:numb_test, :], [-1]),
t_box: test_data['box'][:numb_test, :],
t_type: np.reshape(test_data['type'][:numb_test, :], [-1]),
t_natoms: test_data['natoms_vec'],
t_mesh: test_data['default_mesh'],
is_training: False
}
sess = tf.Session()
sess.run(tf.global_variables_initializer())
[p] = sess.run([polar], feed_dict=feed_dict_test)
p = p.reshape([-1])
refp = [
3.39695248e+01, 2.16564043e+01, 8.18501479e-01, 2.16564043e+01,
1.38211789e+01, 5.22775159e-01, 8.18501479e-01, 5.22775159e-01,
1.97847218e-02, 8.08467431e-01, 3.42081126e+00, -2.01072261e-01,
3.42081126e+00, 1.54924596e+01, -9.06153697e-01, -2.01072261e-01,
-9.06153697e-01, 5.30193262e-02
]
places = 6
for ii in range(p.size):
self.assertAlmostEqual(p[ii], refp[ii], places=places)
def test_model(self):
jfile = 'water_se_r.json'
with open(jfile) as fp:
jdata = json.load (fp)
run_opt = RunOptions(None)
systems = j_must_have(jdata, 'systems')
set_pfx = j_must_have(jdata, 'set_prefix')
batch_size = j_must_have(jdata, 'batch_size')
test_size = j_must_have(jdata, 'numb_test')
batch_size = 1
test_size = 1
stop_batch = j_must_have(jdata, 'stop_batch')
rcut = j_must_have (jdata['model']['descriptor'], 'rcut')
data = DataSystem(systems, set_pfx, batch_size, test_size, rcut, run_opt = None)
test_data = data.get_test ()
numb_test = 1
descrpt = DescrptSeR(jdata['model']['descriptor'])
fitting = EnerFitting(jdata['model']['fitting_net'], descrpt)
model = Model(jdata['model'], descrpt, fitting)
# model._compute_dstats([test_data['coord']], [test_data['box']], [test_data['type']], [test_data['natoms_vec']], [test_data['default_mesh']])
input_data = {'coord' : [test_data['coord']],
'box': [test_data['box']],
'type': [test_data['type']],
'natoms_vec' : [test_data['natoms_vec']],
'default_mesh' : [test_data['default_mesh']]
}
model._compute_input_stat(input_data)
model.descrpt.bias_atom_e = data.compute_energy_shift()
t_prop_c = tf.placeholder(tf.float32, [5], name='t_prop_c')
t_energy = tf.placeholder(global_ener_float_precision, [None], name='t_energy')
t_force = tf.placeholder(global_tf_float_precision, [None], name='t_force')
t_virial = tf.placeholder(global_tf_float_precision, [None], name='t_virial')
t_atom_ener = tf.placeholder(global_tf_float_precision, [None], name='t_atom_ener')
t_coord = tf.placeholder(global_tf_float_precision, [None], name='i_coord')
t_type = tf.placeholder(tf.int32, [None], name='i_type')
t_natoms = tf.placeholder(tf.int32, [model.ntypes+2], name='i_natoms')
t_box = tf.placeholder(global_tf_float_precision, [None, 9], name='i_box')
t_mesh = tf.placeholder(tf.int32, [None], name='i_mesh')
is_training = tf.placeholder(tf.bool)
t_fparam = None
model_pred\
= model.build (t_coord,
t_type,
t_natoms,
t_box,
t_mesh,
t_fparam,
suffix = "se_r",
reuse = False)
energy = model_pred['energy']
force = model_pred['force']
virial = model_pred['virial']
atom_ener = model_pred['atom_ener']
feed_dict_test = {t_prop_c: test_data['prop_c'],
t_energy: test_data['energy'] [:numb_test],
t_force: np.reshape(test_data['force'] [:numb_test, :], [-1]),
t_virial: np.reshape(test_data['virial'] [:numb_test, :], [-1]),
t_atom_ener: np.reshape(test_data['atom_ener'][:numb_test, :], [-1]),
t_coord: np.reshape(test_data['coord'] [:numb_test, :], [-1]),
t_box: test_data['box'] [:numb_test, :],
t_type: np.reshape(test_data['type'] [:numb_test, :], [-1]),
t_natoms: test_data['natoms_vec'],
t_mesh: test_data['default_mesh'],
is_training: False}
sess = tf.Session()
sess.run(tf.global_variables_initializer())
[e, f, v] = sess.run([energy, force, virial],
feed_dict = feed_dict_test)
e = e.reshape([-1])
f = f.reshape([-1])
v = v.reshape([-1])
refe = [6.152085988309423925e+01]
reff = [-1.714443151616400110e-04,-1.315836609370952051e-04,-5.584120460897444674e-06,-7.197863450669731334e-05,-1.384609799994930676e-04,8.856091902774708468e-06,1.120578238869146797e-04,-7.428703645877488470e-05,9.370560731488587317e-07,-1.048347129617610465e-04,1.977876923815685781e-04,7.522050342771599598e-06,2.361772659657814205e-04,-5.774651813388292487e-05,-1.233143271630744828e-05,2.257277740226381951e-08,2.042905031476775584e-04,6.003548585097267914e-07]
refv = [1.035180911513190792e-03,-1.118982949050497126e-04,-2.383287813436022850e-05,-1.118982949050497126e-04,4.362023915782403281e-04,8.119543218224559240e-06,-2.383287813436022850e-05,8.119543218224559240e-06,1.201142938802945237e-06]
refe = np.reshape(refe, [-1])
reff = np.reshape(reff, [-1])
refv = np.reshape(refv, [-1])
places = 6
for ii in range(e.size) :
self.assertAlmostEqual(e[ii], refe[ii], places = places)
for ii in range(f.size) :
self.assertAlmostEqual(f[ii], reff[ii], places = places)
for ii in range(v.size) :
self.assertAlmostEqual(v[ii], refv[ii], places = places)
def test_model(self):
jfile = 'water_se_a_aparam.json'
with open(jfile) as fp:
jdata = json.load(fp)
run_opt = RunOptions(None)
systems = j_must_have(jdata, 'systems')
set_pfx = j_must_have(jdata, 'set_prefix')
batch_size = j_must_have(jdata, 'batch_size')
test_size = j_must_have(jdata, 'numb_test')
batch_size = 1
test_size = 1
stop_batch = j_must_have(jdata, 'stop_batch')
rcut = j_must_have(jdata['model']['descriptor'], 'rcut')
data = DataSystem(systems,
set_pfx,
batch_size,
test_size,
rcut,
run_opt=None)
test_data = data.get_test()
# manually set aparam
test_data['aparam'] = np.load('system/set.000/aparam.npy')
numb_test = 1
descrpt = DescrptSeA(jdata['model']['descriptor'])
fitting = EnerFitting(jdata['model']['fitting_net'], descrpt)
model = Model(jdata['model'], descrpt, fitting)
# model._compute_dstats([test_data['coord']], [test_data['box']], [test_data['type']], [test_data['natoms_vec']], [test_data['default_mesh']])
input_data = {
'coord': [test_data['coord']],
'box': [test_data['box']],
'type': [test_data['type']],
'natoms_vec': [test_data['natoms_vec']],
'default_mesh': [test_data['default_mesh']],
'aparam': [test_data['aparam']],
}
model._compute_dstats(input_data)
model.bias_atom_e = data.compute_energy_shift()
t_prop_c = tf.placeholder(tf.float32, [5], name='t_prop_c')
t_energy = tf.placeholder(global_ener_float_precision, [None],
name='t_energy')
t_force = tf.placeholder(global_tf_float_precision, [None],
name='t_force')
t_virial = tf.placeholder(global_tf_float_precision, [None],
name='t_virial')
t_atom_ener = tf.placeholder(global_tf_float_precision, [None],
name='t_atom_ener')
t_coord = tf.placeholder(global_tf_float_precision, [None],
name='i_coord')
t_type = tf.placeholder(tf.int32, [None], name='i_type')
t_natoms = tf.placeholder(tf.int32, [model.ntypes + 2],
name='i_natoms')
t_box = tf.placeholder(global_tf_float_precision, [None, 9],
name='i_box')
t_mesh = tf.placeholder(tf.int32, [None], name='i_mesh')
t_aparam = tf.placeholder(global_tf_float_precision, [None],
name='i_aparam')
is_training = tf.placeholder(tf.bool)
input_dict = {}
input_dict['aparam'] = t_aparam
model_pred\
= model.build (t_coord,
t_type,
t_natoms,
t_box,
t_mesh,
input_dict,
suffix = "se_a_aparam",
reuse = False)
energy = model_pred['energy']
force = model_pred['force']
virial = model_pred['virial']
atom_ener = model_pred['atom_ener']
feed_dict_test = {
t_prop_c: test_data['prop_c'],
t_energy: test_data['energy'][:numb_test],
t_force: np.reshape(test_data['force'][:numb_test, :], [-1]),
t_virial: np.reshape(test_data['virial'][:numb_test, :], [-1]),
t_atom_ener: np.reshape(test_data['atom_ener'][:numb_test, :],
[-1]),
t_coord: np.reshape(test_data['coord'][:numb_test, :], [-1]),
t_box: test_data['box'][:numb_test, :],
t_type: np.reshape(test_data['type'][:numb_test, :], [-1]),
t_natoms: test_data['natoms_vec'],
t_mesh: test_data['default_mesh'],
t_aparam: np.reshape(test_data['aparam'][:numb_test, :], [-1]),
is_training: False
}
sess = tf.Session()
sess.run(tf.global_variables_initializer())
[e, f, v] = sess.run([energy, force, virial], feed_dict=feed_dict_test)
e = e.reshape([-1])
f = f.reshape([-1])
v = v.reshape([-1])
refe = [61.35473702079649]
reff = [
7.789591210641927388e-02, 9.411176646369459609e-02,
3.785806413688173194e-03, 1.430830954178063386e-01,
1.146964190520970150e-01, -1.320340288927138173e-02,
-7.308720494747594776e-02, 6.508269338140809657e-02,
5.398739145542804643e-04, 5.863268336973800898e-02,
-1.603409523950408699e-01, -5.083084610994957619e-03,
-2.551569799443983988e-01, 3.087934885732580501e-02,
1.508590526622844222e-02, 4.863249399791078065e-02,
-1.444292753594846324e-01, -1.125098094204559241e-03
]
refv = [
-6.069498397488943819e-01, 1.101778888191114192e-01,
1.981907430646132409e-02, 1.101778888191114608e-01,
-3.315612988100872793e-01, -5.999739184898976799e-03,
1.981907430646132756e-02, -5.999739184898974197e-03,
-1.198656608172396325e-03
]
refe = np.reshape(refe, [-1])
reff = np.reshape(reff, [-1])
refv = np.reshape(refv, [-1])
places = 10
for ii in range(e.size):
self.assertAlmostEqual(e[ii], refe[ii], places=places)
for ii in range(f.size):
self.assertAlmostEqual(f[ii], reff[ii], places=places)
for ii in range(v.size):
self.assertAlmostEqual(v[ii], refv[ii], places=places)
def _main():
default_num_inter_threads = 0
parser = argparse.ArgumentParser(description="*** Train a model. ***")
parser.add_argument('INPUT', help='the input json database ')
parser.add_argument(
'-t',
'--inter-threads',
type=int,
default=default_num_inter_threads,
help='With default value %d. ' % default_num_inter_threads +
'Setting the "inter_op_parallelism_threads" key for the tensorflow, ' +
'the "intra_op_parallelism_threads" will be set by the env variable OMP_NUM_THREADS'
)
parser.add_argument(
'--init-model',
type=str,
help='Initialize the model by the provided checkpoint.')
parser.add_argument(
'--restart',
type=str,
help='Restart the training from the provided checkpoint.')
args = parser.parse_args()
# load json database
fp = open(args.INPUT, 'r')
jdata = json.load(fp)
# Setup cluster for distributed training
ps_num = j_must_have(jdata, 'ps_num')
cluster, my_job_name, my_task_index = tf_config_from_slurm(
ps_number=ps_num)
cluster_spec = tf.train.ClusterSpec(cluster)
server = tf.train.Server(server_or_cluster_def=cluster_spec,
job_name=my_job_name,
task_index=my_task_index)
if my_job_name == "ps":
queue = create_done_queue(cluster_spec, my_task_index)
print("create queue")
wait_done_queue(cluster_spec, server, queue, my_task_index)
#server.join()
elif my_job_name == "worker":
is_chief = (my_task_index == 0)
done_ops = connect_done_queue(cluster_spec, my_task_index)
# init params and run options
systems = j_must_have(jdata, 'systems')
set_pfx = j_must_have(jdata, 'set_prefix')
numb_sys = len(systems)
seed = None
if 'seed' in jdata.keys(): seed = jdata['seed']
batch_size = j_must_have(jdata, 'batch_size')
test_size = j_must_have(jdata, 'numb_test')
stop_batch = j_must_have(jdata, 'stop_batch')
rcut = j_must_have(jdata, 'rcut')
data = DataSystem(systems, set_pfx, batch_size, test_size, rcut)
tot_numb_batches = sum(data.get_nbatches())
lr = LearingRate(jdata, tot_numb_batches)
final_lr = lr.value(stop_batch)
run_opt = RunOptions(args)
if is_chief:
print("#")
print("# find %d system(s): " % numb_sys)
print("#")
print(
"# run with intra_op_parallelism_threads = %d, inter_op_parallelism_threads = %d "
% (run_opt.num_intra_threads, run_opt.num_inter_threads))
run_opt.cluster = cluster_spec
run_opt.server = server
run_opt.is_chief = is_chief
run_opt.my_job_name = my_job_name
run_opt.my_task_index = my_task_index
# init the model
model = NNPModel(jdata, run_opt=run_opt)
# build the model with stats from the first system
model.build(data, lr)
start_time = time.time()
cur_batch = 0
if is_chief:
print("# start training, start lr is %e, final lr will be %e" %
(lr.value(cur_batch), final_lr))
sys.stdout.flush()
#model.print_head()
# train the model with the provided systems in a cyclic way
model.train(data, stop_batch)
end_time = time.time()
if is_chief:
print("# finished training")
print("# running time: %.3f s" % (end_time - start_time))
fill_done_queue(cluster_spec, server, done_ops, my_task_index)
def test_model(self):
jfile = 'water_se_a_fparam.json'
with open(jfile) as fp:
jdata = json.load(fp)
run_opt = RunOptions(None)
systems = j_must_have(jdata, 'systems')
set_pfx = j_must_have(jdata, 'set_prefix')
batch_size = j_must_have(jdata, 'batch_size')
test_size = j_must_have(jdata, 'numb_test')
batch_size = 1
test_size = 1
stop_batch = j_must_have(jdata, 'stop_batch')
rcut = j_must_have(jdata['model']['descriptor'], 'rcut')
data = DataSystem(systems,
set_pfx,
batch_size,
test_size,
rcut,
run_opt=None)
test_data = data.get_test()
numb_test = 1
descrpt = DescrptSeA(jdata['model']['descriptor'])
fitting = EnerFitting(jdata['model']['fitting_net'], descrpt)
model = Model(jdata['model'], descrpt, fitting)
model._compute_dstats([test_data['coord']], [test_data['box']],
[test_data['type']], [test_data['natoms_vec']],
[test_data['default_mesh']])
model.bias_atom_e = data.compute_energy_shift()
t_prop_c = tf.placeholder(tf.float32, [5], name='t_prop_c')
t_energy = tf.placeholder(global_ener_float_precision, [None],
name='t_energy')
t_force = tf.placeholder(global_tf_float_precision, [None],
name='t_force')
t_virial = tf.placeholder(global_tf_float_precision, [None],
name='t_virial')
t_atom_ener = tf.placeholder(global_tf_float_precision, [None],
name='t_atom_ener')
t_coord = tf.placeholder(global_tf_float_precision, [None],
name='i_coord')
t_type = tf.placeholder(tf.int32, [None], name='i_type')
t_natoms = tf.placeholder(tf.int32, [model.ntypes + 2],
name='i_natoms')
t_box = tf.placeholder(global_tf_float_precision, [None, 9],
name='i_box')
t_mesh = tf.placeholder(tf.int32, [None], name='i_mesh')
t_fparam = tf.placeholder(global_tf_float_precision, [None],
name='i_fparam')
is_training = tf.placeholder(tf.bool)
input_dict = {}
input_dict['fparam'] = t_fparam
model_pred\
= model.build (t_coord,
t_type,
t_natoms,
t_box,
t_mesh,
input_dict,
suffix = "se_a_fparam",
reuse = False)
energy = model_pred['energy']
force = model_pred['force']
virial = model_pred['virial']
atom_ener = model_pred['atom_ener']
feed_dict_test = {
t_prop_c: test_data['prop_c'],
t_energy: test_data['energy'][:numb_test],
t_force: np.reshape(test_data['force'][:numb_test, :], [-1]),
t_virial: np.reshape(test_data['virial'][:numb_test, :], [-1]),
t_atom_ener: np.reshape(test_data['atom_ener'][:numb_test, :],
[-1]),
t_coord: np.reshape(test_data['coord'][:numb_test, :], [-1]),
t_box: test_data['box'][:numb_test, :],
t_type: np.reshape(test_data['type'][:numb_test, :], [-1]),
t_natoms: test_data['natoms_vec'],
t_mesh: test_data['default_mesh'],
t_fparam: np.reshape(test_data['fparam'][:numb_test, :], [-1]),
is_training: False
}
sess = tf.Session()
sess.run(tf.global_variables_initializer())
[e, f, v] = sess.run([energy, force, virial], feed_dict=feed_dict_test)
e = e.reshape([-1])
f = f.reshape([-1])
v = v.reshape([-1])
refe = [6.135136929183754972e+01]
reff = [
7.761477777656561328e-02, 9.383013575207051205e-02,
3.776776376267230399e-03, 1.428268971463224069e-01,
1.143858253900619654e-01, -1.318441687719179231e-02,
-7.271897092708884403e-02, 6.494907553857684479e-02,
5.355599592111062821e-04, 5.840910251709752199e-02,
-1.599042555763417750e-01, -5.067165555590445389e-03,
-2.546246315216804113e-01, 3.073296814647456451e-02,
1.505994759166155023e-02, 4.849282500878367153e-02,
-1.439937492508420736e-01, -1.120701494357654411e-03
]
refv = [
-6.054303146013112480e-01, 1.097859194719944115e-01,
1.977605183964963390e-02, 1.097859194719943976e-01,
-3.306167096812382966e-01, -5.978855662865613894e-03,
1.977605183964964083e-02, -5.978855662865616497e-03,
-1.196331922996723236e-03
]
refe = np.reshape(refe, [-1])
reff = np.reshape(reff, [-1])
refv = np.reshape(refv, [-1])
places = 10
for ii in range(e.size):
self.assertAlmostEqual(e[ii], refe[ii], places=places)
for ii in range(f.size):
self.assertAlmostEqual(f[ii], reff[ii], places=places)
for ii in range(v.size):
self.assertAlmostEqual(v[ii], refv[ii], places=places)