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)