def test_py_interface(self): hh = 1e-4 places = 4 sess = tf.Session() t_energy, t_force, t_virial \ = op_module.ewald_recp(self.coord, self.charge, self.nloc, self.box, ewald_h = self.ewald_h, ewald_beta = self.ewald_beta) [e, f, v] = sess.run([t_energy, t_force, t_virial], feed_dict = { self.coord: self.dcoord.reshape([-1]), self.charge: self.dcharge.reshape([-1]), self.box: self.dbox.reshape([-1]), self.nloc: [self.natoms], }) er = EwaldRecp(self.ewald_h, self.ewald_beta) e1, f1, v1 = er.eval(self.dcoord, self.dcharge, self.dbox) for ff in range(self.nframes): self.assertAlmostEqual(e[ff], e1[ff], places = places, msg = "frame %d energy failed" % (ff)) for idx in range(self.natoms): for dd in range(3): self.assertAlmostEqual(f[ff, idx*3+dd], f1[ff,idx*3+dd], places = places, msg = "frame %d force component [%d,%d] failed" % (ff, idx, dd)) for d0 in range(3): for d1 in range(3): self.assertAlmostEqual(v[ff, d0*3+d1], v[ff,d0*3+d1], places = places, msg = "frame %d virial component [%d,%d] failed" % (ff, d0, d1))
def test_op_gelu(self): w = tf.constant( [[0.1, 0.2, 0.3, 0.4], [0.5, 0.6, 0.7, 0.8], [0.9, 1, 1.1, 1.2]], dtype='double') x = tf.constant([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9], [1.0, 1.1, 1.2]], dtype='double') b = tf.constant([[0.1], [0.2], [0.3], [0.4]], dtype='double') xbar = tf.matmul(x, w) + b y = gelu(xbar) dy = op_module.unaggregated_dy_dx_s(y, w, xbar, tf.constant(2)) dy_array = tf.Session().run(dy) answer = np.array( [[ 8.549286163555620821e-02, 1.782905778685600906e-01, 2.776474599997448833e-01, 3.827650237273348965e-01 ], [ 1.089906023807040714e-01, 2.230820937721638697e-01, 3.381867859682909927e-01, 4.513008399758057232e-01 ], [ 1.124254240556722684e-01, 2.209918074710395253e-01, 3.238894323148118759e-01, 4.220357318198978414e-01 ], [ 1.072173273655498138e-01, 2.082159073100979807e-01, 3.059816075270163083e-01, 4.032981557798429595e-01 ]]) places = 18 np.testing.assert_almost_equal(dy_array, answer, places)
def get_tensor_by_name_from_graph(graph: tf.Graph, tensor_name: str) -> tf.Tensor: """ Load tensor value from the given tf.Graph object Parameters ---------- graph : tf.Graph The input TensorFlow graph tensor_name : str Indicates which tensor which will be loaded from the frozen model Returns ------- tf.Tensor The tensor which was loaded from the frozen model Raises ------ GraphWithoutTensorError Whether the tensor_name is within the frozen model """ try: tensor = graph.get_tensor_by_name(tensor_name + ':0') except KeyError as e: raise GraphWithoutTensorError() from e with tf.Session(graph=graph) as sess: tensor = run_sess(sess, tensor) return tensor
def __init__(self, model_file): model_file = model_file graph = self.load_graph(model_file) t_mt = graph.get_tensor_by_name('load/model_attr/model_type:0') sess = tf.Session(graph=graph) [mt] = sess.run([t_mt], feed_dict={}) self.model_type = mt.decode('utf-8')
def test_op_tanh(self): w = tf.constant( [[0.1, 0.2, 0.3, 0.4], [0.5, 0.6, 0.7, 0.8], [0.9, 1, 1.1, 1.2]], dtype='double') x = tf.constant([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9], [1.0, 1.1, 1.2]], dtype='double') b = tf.constant([[0.1], [0.2], [0.3], [0.4]], dtype='double') xbar = tf.matmul(x, w) + b y = tf.nn.tanh(xbar) dy = op_module.unaggregated_dy_dx_s(y, w, xbar, tf.constant(1)) dy_array = tf.Session().run(dy) answer = np.array( [[ 8.008666403121351973e-02, 1.513925729426658651e-01, 2.134733287761668430e-01, 2.661983049806041501e-01 ], [ 4.010658815015744061e-02, 6.306476628799793926e-02, 7.332167904608145881e-02, 7.494218676568849269e-02 ], [ 1.561705624394135218e-02, 1.994112926507514427e-02, 1.887519955881525671e-02, 1.576442161040989692e-02 ], [ 5.492686739421748753e-03, 5.754985286040992763e-03, 4.493113544969218158e-03, 3.107638130764600777e-03 ]]) places = 18 np.testing.assert_almost_equal(dy_array, answer, places)
def _init_sess_serial(self) : self.sess = tf.Session( config=tf.ConfigProto(intra_op_parallelism_threads=self.run_opt.num_intra_threads, inter_op_parallelism_threads=self.run_opt.num_inter_threads )) self.saver = tf.train.Saver() saver = self.saver if self.run_opt.init_mode == 'init_from_scratch' : self._message("initialize model from scratch") init_op = tf.global_variables_initializer() self.sess.run(init_op) fp = open(self.disp_file, "w") fp.close () elif self.run_opt.init_mode == 'init_from_model' : self._message("initialize from model %s" % self.run_opt.init_model) init_op = tf.global_variables_initializer() self.sess.run(init_op) saver.restore (self.sess, self.run_opt.init_model) self.sess.run(self.global_step.assign(0)) fp = open(self.disp_file, "w") fp.close () elif self.run_opt.init_mode == 'restart' : self._message("restart from model %s" % self.run_opt.restart) init_op = tf.global_variables_initializer() self.sess.run(init_op) saver.restore (self.sess, self.run_opt.restart) else : raise RuntimeError ("unkown init mode")
def __init__(self, model_file, variable_name, variable_dof) : DeepEval.__init__(self, model_file) self.model_file = model_file self.graph = self.load_graph (self.model_file) self.variable_name = variable_name self.variable_dof = variable_dof # checkout input/output tensors from graph self.t_ntypes = self.graph.get_tensor_by_name ('load/descrpt_attr/ntypes:0') self.t_rcut = self.graph.get_tensor_by_name ('load/descrpt_attr/rcut:0') self.t_tmap = self.graph.get_tensor_by_name ('load/model_attr/tmap:0') self.t_sel_type= self.graph.get_tensor_by_name ('load/model_attr/sel_type:0') # inputs self.t_coord = self.graph.get_tensor_by_name ('load/t_coord:0') self.t_type = self.graph.get_tensor_by_name ('load/t_type:0') self.t_natoms = self.graph.get_tensor_by_name ('load/t_natoms:0') self.t_box = self.graph.get_tensor_by_name ('load/t_box:0') self.t_mesh = self.graph.get_tensor_by_name ('load/t_mesh:0') # outputs self.t_tensor = self.graph.get_tensor_by_name ('load/o_%s:0' % self.variable_name) # start a tf session associated to the graph self.sess = tf.Session (graph = self.graph) [self.ntypes, self.rcut, self.tmap, self.tselt] = self.sess.run([self.t_ntypes, self.t_rcut, self.t_tmap, self.t_sel_type]) self.tmap = self.tmap.decode('UTF-8').split()
def _compute_dstats_sys_smth (self, data_coord, data_box, data_atype, natoms_vec, mesh) : avg_zero = np.zeros([self.ntypes,self.ndescrpt]).astype(global_np_float_precision) std_ones = np.ones ([self.ntypes,self.ndescrpt]).astype(global_np_float_precision) sub_graph = tf.Graph() with sub_graph.as_default(): descrpt, descrpt_deriv, rij, nlist \ = op_module.descrpt_se_a (tf.constant(data_coord), tf.constant(data_atype), tf.constant(natoms_vec, dtype = tf.int32), tf.constant(data_box), tf.constant(mesh), tf.constant(avg_zero), tf.constant(std_ones), rcut_a = self.rcut_a, rcut_r = self.rcut_r, rcut_r_smth = self.rcut_r_smth, sel_a = self.sel_a, sel_r = self.sel_r) # self.sess.run(tf.global_variables_initializer()) # sub_sess = tf.Session(graph = sub_graph, # config=tf.ConfigProto(intra_op_parallelism_threads=self.run_opt.num_intra_threads, # inter_op_parallelism_threads=self.run_opt.num_inter_threads # )) sub_sess = tf.Session(graph = sub_graph) dd_all = sub_sess.run(descrpt) sub_sess.close() natoms = natoms_vec dd_all = np.reshape(dd_all, [-1, self.ndescrpt * natoms[0]]) start_index = 0 sysr = [] sysa = [] sysn = [] sysr2 = [] sysa2 = [] for type_i in range(self.ntypes): end_index = start_index + self.ndescrpt * natoms[2+type_i] dd = dd_all[:, start_index:end_index] dd = np.reshape(dd, [-1, self.ndescrpt]) start_index = end_index # compute dd = np.reshape (dd, [-1, 4]) ddr = dd[:,:1] dda = dd[:,1:] sumr = np.sum(ddr) suma = np.sum(dda) / 3. sumn = dd.shape[0] sumr2 = np.sum(np.multiply(ddr, ddr)) suma2 = np.sum(np.multiply(dda, dda)) / 3. sysr.append(sumr) sysa.append(suma) sysn.append(sumn) sysr2.append(sumr2) sysa2.append(suma2) return sysr, sysr2, sysa, sysa2, sysn
def __init__(self, data): self.sess = tf.Session() self.data = data self.natoms = self.data.get_natoms() self.ntypes = self.data.get_ntypes() self.sel = [12, 24] self.sel_a = [0, 0] self.rcut_smth = 2.45 self.rcut = 10.0 self.nnei = np.cumsum(self.sel)[-1] self.ndescrpt = self.nnei * 1 davg = np.zeros([self.ntypes, self.ndescrpt]) dstd = np.ones([self.ntypes, self.ndescrpt]) self.t_avg = tf.constant(davg.astype(global_np_float_precision)) self.t_std = tf.constant(dstd.astype(global_np_float_precision)) self.default_mesh = np.zeros(6, dtype=np.int32) self.default_mesh[3] = 2 self.default_mesh[4] = 2 self.default_mesh[5] = 2 # make place holder self.coord = tf.placeholder(global_tf_float_precision, [None, self.natoms[0] * 3], name='t_coord') self.box = tf.placeholder(global_tf_float_precision, [None, 9], name='t_box') self.type = tf.placeholder(tf.int32, [None, self.natoms[0]], name="t_type") self.tnatoms = tf.placeholder(tf.int32, [None], name="t_natoms")
def __init__(self, model_file): self.model_file = model_file self.graph = self.load_graph(self.model_file) # checkout input/output tensors from graph self.t_ntypes = self.graph.get_tensor_by_name( 'load/descrpt_attr/ntypes:0') self.t_rcut = self.graph.get_tensor_by_name('load/descrpt_attr/rcut:0') self.t_dfparam = self.graph.get_tensor_by_name( 'load/fitting_attr/dfparam:0') self.t_tmap = self.graph.get_tensor_by_name('load/model_attr/tmap:0') # inputs self.t_coord = self.graph.get_tensor_by_name('load/t_coord:0') self.t_type = self.graph.get_tensor_by_name('load/t_type:0') self.t_natoms = self.graph.get_tensor_by_name('load/t_natoms:0') self.t_box = self.graph.get_tensor_by_name('load/t_box:0') self.t_mesh = self.graph.get_tensor_by_name('load/t_mesh:0') # outputs self.t_energy = self.graph.get_tensor_by_name('load/o_energy:0') self.t_force = self.graph.get_tensor_by_name('load/o_force:0') self.t_virial = self.graph.get_tensor_by_name('load/o_virial:0') self.t_ae = self.graph.get_tensor_by_name('load/o_atom_energy:0') self.t_av = self.graph.get_tensor_by_name('load/o_atom_virial:0') self.t_fparam = None # check if the graph has fparam for op in self.graph.get_operations(): if op.name == 'load/t_fparam': self.t_fparam = self.graph.get_tensor_by_name( 'load/t_fparam:0') self.has_fparam = self.t_fparam is not None # start a tf session associated to the graph self.sess = tf.Session(graph=self.graph) [self.ntypes, self.rcut, self.dfparam, self.tmap] = self.sess.run( [self.t_ntypes, self.t_rcut, self.t_dfparam, self.t_tmap]) self.tmap = self.tmap.decode('UTF-8').split()
def __init__(self, ntypes: int, rcut: float) -> None: """ Constructor """ self.rcut = rcut self.ntypes = ntypes self.place_holders = {} sub_graph = tf.Graph() with sub_graph.as_default(): for ii in ['coord', 'box']: self.place_holders[ii] = tf.placeholder( GLOBAL_NP_FLOAT_PRECISION, [None, None], name='t_' + ii) self.place_holders['type'] = tf.placeholder(tf.int32, [None, None], name='t_type') self.place_holders['natoms_vec'] = tf.placeholder( tf.int32, [self.ntypes + 2], name='t_natoms') self.place_holders['default_mesh'] = tf.placeholder(tf.int32, [None], name='t_mesh') self._max_nbor_size, self._min_nbor_dist \ = op_module.neighbor_stat(self.place_holders['coord'], self.place_holders['type'], self.place_holders['natoms_vec'], self.place_holders['box'], self.place_holders['default_mesh'], rcut = self.rcut) self.sub_sess = tf.Session(graph=sub_graph, config=default_tf_session_config)
def __init__(self, rcut: float, sel_a : List[int], sel_r : List[int], axis_rule : List[int] ) -> None: """ Constructor """ # args = ClassArg()\ # .add('sel_a', list, must = True) \ # .add('sel_r', list, must = True) \ # .add('rcut', float, default = 6.0) \ # .add('axis_rule',list, must = True) # class_data = args.parse(jdata) self.sel_a = sel_a self.sel_r = sel_r self.axis_rule = axis_rule self.rcut_r = rcut # ntypes and rcut_a === -1 self.ntypes = len(self.sel_a) assert(self.ntypes == len(self.sel_r)) self.rcut_a = -1 # 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 self.davg = None self.dstd = None self.place_holders = {} avg_zero = np.zeros([self.ntypes,self.ndescrpt]).astype(GLOBAL_NP_FLOAT_PRECISION) std_ones = np.ones ([self.ntypes,self.ndescrpt]).astype(GLOBAL_NP_FLOAT_PRECISION) sub_graph = tf.Graph() with sub_graph.as_default(): name_pfx = 'd_lf_' for ii in ['coord', 'box']: self.place_holders[ii] = tf.placeholder(GLOBAL_NP_FLOAT_PRECISION, [None, None], name = name_pfx+'t_'+ii) self.place_holders['type'] = tf.placeholder(tf.int32, [None, None], name=name_pfx+'t_type') self.place_holders['natoms_vec'] = tf.placeholder(tf.int32, [self.ntypes+2], name=name_pfx+'t_natoms') self.place_holders['default_mesh'] = tf.placeholder(tf.int32, [None], name=name_pfx+'t_mesh') self.stat_descrpt, descrpt_deriv, rij, nlist, axis, rot_mat \ = op_module.descrpt (self.place_holders['coord'], self.place_holders['type'], self.place_holders['natoms_vec'], self.place_holders['box'], self.place_holders['default_mesh'], tf.constant(avg_zero), tf.constant(std_ones), rcut_a = self.rcut_a, rcut_r = self.rcut_r, sel_a = self.sel_a, sel_r = self.sel_r, axis_rule = self.axis_rule) self.sub_sess = tf.Session(graph = sub_graph, config=default_tf_session_config)
def __init__(self, model_file, default_tf_graph = False) : DeepEval.__init__(self, model_file, default_tf_graph = default_tf_graph) # self.model_file = model_file # self.graph = self.load_graph (self.model_file) # checkout input/output tensors from graph self.t_ntypes = self.graph.get_tensor_by_name ('load/descrpt_attr/ntypes:0') self.t_rcut = self.graph.get_tensor_by_name ('load/descrpt_attr/rcut:0') self.t_dfparam= self.graph.get_tensor_by_name ('load/fitting_attr/dfparam:0') self.t_daparam= self.graph.get_tensor_by_name ('load/fitting_attr/daparam:0') self.t_tmap = self.graph.get_tensor_by_name ('load/model_attr/tmap:0') # inputs self.t_coord = self.graph.get_tensor_by_name ('load/t_coord:0') self.t_type = self.graph.get_tensor_by_name ('load/t_type:0') self.t_natoms = self.graph.get_tensor_by_name ('load/t_natoms:0') self.t_box = self.graph.get_tensor_by_name ('load/t_box:0') self.t_mesh = self.graph.get_tensor_by_name ('load/t_mesh:0') # outputs self.t_energy = self.graph.get_tensor_by_name ('load/o_energy:0') self.t_force = self.graph.get_tensor_by_name ('load/o_force:0') self.t_virial = self.graph.get_tensor_by_name ('load/o_virial:0') self.t_ae = self.graph.get_tensor_by_name ('load/o_atom_energy:0') self.t_av = self.graph.get_tensor_by_name ('load/o_atom_virial:0') self.t_fparam = None self.t_aparam = None # check if the graph has fparam for op in self.graph.get_operations(): if op.name == 'load/t_fparam' : self.t_fparam = self.graph.get_tensor_by_name ('load/t_fparam:0') self.has_fparam = self.t_fparam is not None # check if the graph has aparam for op in self.graph.get_operations(): if op.name == 'load/t_aparam' : self.t_aparam = self.graph.get_tensor_by_name ('load/t_aparam:0') self.has_aparam = self.t_aparam is not None # start a tf session associated to the graph self.sess = tf.Session (graph = self.graph, config=default_tf_session_config) [self.ntypes, self.rcut, self.dfparam, self.daparam, self.tmap] = self.sess.run([self.t_ntypes, self.t_rcut, self.t_dfparam, self.t_daparam, self.t_tmap]) self.tmap = self.tmap.decode('UTF-8').split() # setup modifier try: t_modifier_type = self.graph.get_tensor_by_name('load/modifier_attr/type:0') self.modifier_type = self.sess.run(t_modifier_type).decode('UTF-8') except ValueError: self.modifier_type = None except KeyError: self.modifier_type = None if self.modifier_type == 'dipole_charge': t_mdl_name = self.graph.get_tensor_by_name('load/modifier_attr/mdl_name:0') t_mdl_charge_map = self.graph.get_tensor_by_name('load/modifier_attr/mdl_charge_map:0') t_sys_charge_map = self.graph.get_tensor_by_name('load/modifier_attr/sys_charge_map:0') t_ewald_h = self.graph.get_tensor_by_name('load/modifier_attr/ewald_h:0') t_ewald_beta = self.graph.get_tensor_by_name('load/modifier_attr/ewald_beta:0') [mdl_name, mdl_charge_map, sys_charge_map, ewald_h, ewald_beta] = self.sess.run([t_mdl_name, t_mdl_charge_map, t_sys_charge_map, t_ewald_h, t_ewald_beta]) mdl_charge_map = [int(ii) for ii in mdl_charge_map.decode('UTF-8').split()] sys_charge_map = [int(ii) for ii in sys_charge_map.decode('UTF-8').split()] self.dm = DipoleChargeModifier(mdl_name, mdl_charge_map, sys_charge_map, ewald_h = ewald_h, ewald_beta = ewald_beta)
def __init__(self, model_file, load_prefix='load', default_tf_graph=False): self.graph = self._load_graph(model_file, prefix=load_prefix, default_tf_graph=default_tf_graph) t_mt = self.graph.get_tensor_by_name( os.path.join(load_prefix, 'model_attr/model_type:0')) sess = tf.Session(graph=self.graph, config=default_tf_session_config) [mt] = sess.run([t_mt], feed_dict={}) self.model_type = mt.decode('utf-8')
def model_type(self) -> str: """Get type of model. :type:str """ if not self._model_type: t_mt = self._get_tensor("model_attr/model_type:0") sess = tf.Session(graph=self.graph, config=default_tf_session_config) [mt] = run_sess(sess, [t_mt], feed_dict={}) self._model_type = mt.decode("utf-8") return self._model_type
def freeze_graph(model_folder, output, output_node_names=None): # We retrieve our checkpoint fullpath checkpoint = tf.train.get_checkpoint_state(model_folder) input_checkpoint = checkpoint.model_checkpoint_path # We precise the file fullname of our freezed graph absolute_model_folder = "/".join(input_checkpoint.split('/')[:-1]) output_graph = absolute_model_folder + "/" + output # Before exporting our graph, we need to precise what is our output node # This is how TF decides what part of the Graph he has to keep and what part it can dump # NOTE: this variable is plural, because you can have multiple output nodes # output_node_names = "energy_test,force_test,virial_test,t_rcut" # We clear devices to allow TensorFlow to control on which device it will load operations clear_devices = True # We import the meta graph and retrieve a Saver saver = tf.train.import_meta_graph(input_checkpoint + '.meta', clear_devices=clear_devices) # We retrieve the protobuf graph definition graph = tf.get_default_graph() input_graph_def = graph.as_graph_def() nodes = [n.name for n in input_graph_def.node] # We start a session and restore the graph weights with tf.Session() as sess: saver.restore(sess, input_checkpoint) model_type = sess.run('model_attr/model_type:0', feed_dict={}).decode('utf-8') if 'modifier_attr/type' in nodes: modifier_type = sess.run('modifier_attr/type:0', feed_dict={}).decode('utf-8') else: modifier_type = None if output_node_names is None: output_node_names = _make_node_names(model_type, modifier_type) print('The following nodes will be frozen: %s' % output_node_names) # We use a built-in TF helper to export variables to constants output_graph_def = tf.graph_util.convert_variables_to_constants( sess, # The session is used to retrieve the weights input_graph_def, # The graph_def is used to retrieve the nodes output_node_names.split( "," ) # The output node names are used to select the usefull nodes ) # Finally we serialize and dump the output graph to the filesystem with tf.gfile.GFile(output_graph, "wb") as f: f.write(output_graph_def.SerializeToString()) print("%d ops in the final graph." % len(output_graph_def.node))
def _compute_dstats_sys_nonsmth (self, data_coord, data_box, data_atype, natoms_vec, mesh, reuse = None) : avg_zero = np.zeros([self.ntypes,self.ndescrpt]).astype(global_np_float_precision) std_ones = np.ones ([self.ntypes,self.ndescrpt]).astype(global_np_float_precision) sub_graph = tf.Graph() with sub_graph.as_default(): descrpt, descrpt_deriv, rij, nlist, axis, rot_mat \ = op_module.descrpt (tf.constant(data_coord), tf.constant(data_atype), tf.constant(natoms_vec, dtype = tf.int32), tf.constant(data_box), tf.constant(mesh), tf.constant(avg_zero), tf.constant(std_ones), rcut_a = self.rcut_a, rcut_r = self.rcut_r, sel_a = self.sel_a, sel_r = self.sel_r, axis_rule = self.axis_rule) # self.sess.run(tf.global_variables_initializer()) # sub_sess = tf.Session(graph = sub_graph, # config=tf.ConfigProto(intra_op_parallelism_threads=self.run_opt.num_intra_threads, # inter_op_parallelism_threads=self.run_opt.num_inter_threads # )) sub_sess = tf.Session(graph = sub_graph) dd_all = sub_sess.run(descrpt) sub_sess.close() natoms = natoms_vec dd_all = np.reshape(dd_all, [-1, self.ndescrpt * natoms[0]]) start_index = 0 sysv = [] sysn = [] sysv2 = [] for type_i in range(self.ntypes): end_index = start_index + self.ndescrpt * natoms[2+type_i] dd = dd_all[:, start_index:end_index] dd = np.reshape(dd, [-1, self.ndescrpt]) start_index = end_index # compute sumv = np.sum(dd, axis = 0) sumn = dd.shape[0] sumv2 = np.sum(np.multiply(dd,dd), axis = 0) sysv.append(sumv) sysn.append(sumn) sysv2.append(sumv2) return sysv, sysv2, sysn
def _init_session(self): config = get_tf_session_config() device, idx = self.run_opt.my_device.split(":", 1) if device == "gpu": config.gpu_options.visible_device_list = idx self.sess = tf.Session(config=config) # Initializes or restore global variables init_op = tf.global_variables_initializer() if self.run_opt.is_chief: self.saver = tf.train.Saver(save_relative_paths=True) if self.run_opt.init_mode == 'init_from_scratch' : log.info("initialize model from scratch") run_sess(self.sess, init_op) if not self.is_compress: fp = open(self.disp_file, "w") fp.close () elif self.run_opt.init_mode == 'init_from_model' : log.info("initialize from model %s" % self.run_opt.init_model) run_sess(self.sess, init_op) self.saver.restore (self.sess, self.run_opt.init_model) run_sess(self.sess, self.global_step.assign(0)) fp = open(self.disp_file, "w") fp.close () elif self.run_opt.init_mode == 'restart' : log.info("restart from model %s" % self.run_opt.restart) run_sess(self.sess, init_op) self.saver.restore (self.sess, self.run_opt.restart) elif self.run_opt.init_mode == 'init_from_frz_model' : log.info("initialize training from the frozen model") run_sess(self.sess, init_op) fp = open(self.disp_file, "w") fp.close () else : raise RuntimeError ("unkown init mode") else: run_sess(self.sess, init_op) self.saver = None # Ensure variable consistency among tasks when training starts if self.run_opt.is_distrib: bcast_op = self.run_opt._HVD.broadcast_global_variables(0) if self.run_opt.is_chief: log.info('broadcast global variables to other tasks') else: log.info('receive global variables from task#0') run_sess(self.sess, bcast_op)
def setUp(self, data): self.sess = tf.Session() self.data = data self.natoms = self.data.get_natoms() self.ntypes = self.data.get_ntypes() param_a = { 'sel': [12, 24], 'rcut': 4, 'rcut_smth': 3.5, 'neuron': [5, 10, 20], 'seed': 1, } param_r = { 'sel': [20, 40], 'rcut': 6, 'rcut_smth': 6.5, 'neuron': [10, 20, 40], 'seed': 1, } param = {'a': param_a, 'r': param_r} self.descrpt = DescrptSeAR(param) self.ndescrpt = self.descrpt.get_dim_out() # davg = np.zeros ([self.ntypes, self.ndescrpt]) # dstd = np.ones ([self.ntypes, self.ndescrpt]) # self.t_avg = tf.constant(davg.astype(np.float64)) # self.t_std = tf.constant(dstd.astype(np.float64)) avg_a = np.zeros([self.ntypes, self.descrpt.descrpt_a.ndescrpt]) std_a = np.ones([self.ntypes, self.descrpt.descrpt_a.ndescrpt]) avg_r = np.zeros([self.ntypes, self.descrpt.descrpt_r.ndescrpt]) std_r = np.ones([self.ntypes, self.descrpt.descrpt_r.ndescrpt]) self.avg = [avg_a, avg_r] self.std = [std_a, std_r] self.default_mesh = np.zeros(6, dtype=np.int32) self.default_mesh[3] = 2 self.default_mesh[4] = 2 self.default_mesh[5] = 2 # make place holder self.coord = tf.placeholder(global_tf_float_precision, [None, self.natoms[0] * 3], name='t_coord') self.box = tf.placeholder(global_tf_float_precision, [None, 9], name='t_box') self.type = tf.placeholder(tf.int32, [None, self.natoms[0]], name="t_type") self.tnatoms = tf.placeholder(tf.int32, [None], name="t_natoms")
def __init__(self, hh, beta): self.hh = hh self.beta = beta with tf.Graph().as_default() as graph: # place holders self.t_nloc = tf.placeholder(tf.int32, [1], name="t_nloc") self.t_coord = tf.placeholder(global_tf_float_precision, [None], name='t_coord') self.t_charge = tf.placeholder(global_tf_float_precision, [None], name='t_charge') self.t_box = tf.placeholder(global_tf_float_precision, [None], name='t_box') # output self.t_energy, self.t_force, self.t_virial \ = op_module.ewald_recp(self.t_coord, self.t_charge, self.t_nloc, self.t_box, ewald_h = self.hh, ewald_beta = self.beta) self.sess = tf.Session(graph=graph, config=default_tf_session_config)
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, model_file, variable_name, variable_dof, load_prefix='load', default_tf_graph=False): DeepEval.__init__(self, model_file, load_prefix=load_prefix, default_tf_graph=default_tf_graph) # self.model_file = model_file # self.graph = self.load_graph (self.model_file) self.variable_name = variable_name self.variable_dof = variable_dof # checkout input/output tensors from graph self.t_ntypes = self.graph.get_tensor_by_name( os.path.join(load_prefix, 'descrpt_attr/ntypes:0')) self.t_rcut = self.graph.get_tensor_by_name( os.path.join(load_prefix, 'descrpt_attr/rcut:0')) self.t_tmap = self.graph.get_tensor_by_name( os.path.join(load_prefix, 'model_attr/tmap:0')) self.t_sel_type = self.graph.get_tensor_by_name( os.path.join(load_prefix, 'model_attr/sel_type:0')) # inputs self.t_coord = self.graph.get_tensor_by_name( os.path.join(load_prefix, 't_coord:0')) self.t_type = self.graph.get_tensor_by_name( os.path.join(load_prefix, 't_type:0')) self.t_natoms = self.graph.get_tensor_by_name( os.path.join(load_prefix, 't_natoms:0')) self.t_box = self.graph.get_tensor_by_name( os.path.join(load_prefix, 't_box:0')) self.t_mesh = self.graph.get_tensor_by_name( os.path.join(load_prefix, 't_mesh:0')) # outputs self.t_tensor = self.graph.get_tensor_by_name( os.path.join(load_prefix, 'o_%s:0' % self.variable_name)) # start a tf session associated to the graph self.sess = tf.Session(graph=self.graph, config=default_tf_session_config) [self.ntypes, self.rcut, self.tmap, self.tselt] = self.sess.run( [self.t_ntypes, self.t_rcut, self.t_tmap, self.t_sel_type]) self.tmap = self.tmap.decode('UTF-8').split()
def __init__(self, model_file: "Path", load_prefix: str = 'load', default_tf_graph: bool = False) -> None: """Constructor""" DeepEval.__init__(self, model_file, load_prefix=load_prefix, default_tf_graph=default_tf_graph) # check model type model_type = self.tensors["t_tensor"][2:-2] assert self.model_type == model_type, \ f"expect {model_type} model but got {self.model_type}" # now load tensors to object attributes for attr_name, tensor_name in self.tensors.items(): self._get_tensor(tensor_name, attr_name) # load optional tensors if possible optional_tensors = { "t_global_tensor": f"o_global_{model_type}:0", "t_force": "o_force:0", "t_virial": "o_virial:0", "t_atom_virial": "o_atom_virial:0" } try: # first make sure these tensor all exists (but do not modify self attr) for attr_name, tensor_name in optional_tensors.items(): self._get_tensor(tensor_name) # then put those into self.attrs for attr_name, tensor_name in optional_tensors.items(): self._get_tensor(tensor_name, attr_name) except KeyError: self._support_gfv = False else: self.tensors.update(optional_tensors) self._support_gfv = True # start a tf session associated to the graph self.sess = tf.Session(graph=self.graph, config=default_tf_session_config) self._run_default_sess() self.tmap = self.tmap.decode('UTF-8').split()
def model_version(self) -> str: """Get version of model. Returns ------- str version of model """ if not self._model_version: try: t_mt = self._get_tensor("model_attr/model_version:0") except KeyError: # For deepmd-kit version 0.x - 1.x, set model version to 0.0 self._model_version = "0.0" else: sess = tf.Session(graph=self.graph, config=default_tf_session_config) [mt] = run_sess(sess, [t_mt], feed_dict={}) self._model_version = mt.decode("utf-8") return self._model_version
def test_force(self): hh = 1e-4 places = 4 sess = tf.Session() t_energy, t_force, t_virial \ = op_module.ewald_recp(self.coord, self.charge, self.nloc, self.box, ewald_h = self.ewald_h, ewald_beta = self.ewald_beta) [force] = sess.run([t_force], feed_dict = { self.coord: self.dcoord.reshape([-1]), self.charge: self.dcharge.reshape([-1]), self.box: self.dbox.reshape([-1]), self.nloc: [self.natoms], }) for idx in range(self.natoms): for dd in range(3): dcoordp = np.copy(self.dcoord) dcoordm = np.copy(self.dcoord) dcoordp[:,idx*3+dd] = self.dcoord[:,idx*3+dd] + hh dcoordm[:,idx*3+dd] = self.dcoord[:,idx*3+dd] - hh energyp = sess.run([t_energy], feed_dict = { self.coord: dcoordp.reshape([-1]), self.charge: self.dcharge.reshape([-1]), self.box: self.dbox.reshape([-1]), self.nloc: [self.natoms], }) energym = sess.run([t_energy], feed_dict = { self.coord: dcoordm.reshape([-1]), self.charge: self.dcharge.reshape([-1]), self.box: self.dbox.reshape([-1]), self.nloc: [self.natoms], }) c_force = -(energyp[0] - energym[0]) / (2*hh) for ff in range(self.nframes): self.assertAlmostEqual(c_force[ff], force[ff,idx*3+dd], places = places, msg = "frame %d force component [%d,%d] failed" % (ff, idx, dd))
def setUp(self, data, comp=0, pbc=True): self.sess = tf.Session() self.data = data self.natoms = self.data.get_natoms() self.ntypes = self.data.get_ntypes() self.sel_a = [12, 24] self.sel_r = [12, 24] self.rcut_a = -1 self.rcut_r = 10.0 self.axis_rule = [0, 1, 0, 0, 1, 1, 0, 0, 0, 0, 1, 0] 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 davg = np.zeros([self.ntypes, self.ndescrpt]) dstd = np.ones([self.ntypes, self.ndescrpt]) self.t_avg = tf.constant(davg.astype(global_np_float_precision)) self.t_std = tf.constant(dstd.astype(global_np_float_precision)) if pbc: self.default_mesh = np.zeros(6, dtype=np.int32) self.default_mesh[3] = 2 self.default_mesh[4] = 2 self.default_mesh[5] = 2 else: self.default_mesh = np.array([], dtype=np.int32) # make place holder self.coord = tf.placeholder(global_tf_float_precision, [None, self.natoms[0] * 3], name='t_coord') self.box = tf.placeholder(global_tf_float_precision, [None, 9], name='t_box') self.type = tf.placeholder(tf.int32, [None, self.natoms[0]], name="t_type") self.tnatoms = tf.placeholder(tf.int32, [None], name="t_natoms")
def __init__(self, jdata): args = ClassArg()\ .add('sel', list, must = True) \ .add('rcut', float, default = 6.0) \ .add('rcut_smth',float, default = 5.5) \ .add('neuron', list, default = [10, 20, 40]) \ .add('axis_neuron', int, default = 4, alias = 'n_axis_neuron') \ .add('resnet_dt',bool, default = False) \ .add('trainable',bool, default = True) \ .add('seed', int) \ .add('exclude_types', list, default = []) \ .add('set_davg_zero', bool, default = False) \ .add('activation_function', str, default = 'tanh') \ .add('precision', str, default = "default") class_data = args.parse(jdata) self.sel_a = class_data['sel'] self.rcut_r = class_data['rcut'] self.rcut_r_smth = class_data['rcut_smth'] self.filter_neuron = class_data['neuron'] self.n_axis_neuron = class_data['axis_neuron'] self.filter_resnet_dt = class_data['resnet_dt'] self.seed = class_data['seed'] self.trainable = class_data['trainable'] self.filter_activation_fn = get_activation_func( class_data['activation_function']) self.filter_precision = get_precision(class_data['precision']) exclude_types = class_data['exclude_types'] self.exclude_types = set() for tt in exclude_types: assert (len(tt) == 2) self.exclude_types.add((tt[0], tt[1])) self.exclude_types.add((tt[1], tt[0])) self.set_davg_zero = class_data['set_davg_zero'] # descrpt config self.sel_r = [0 for ii in range(len(self.sel_a))] self.ntypes = len(self.sel_a) assert (self.ntypes == len(self.sel_r)) self.rcut_a = -1 # 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 self.useBN = False self.dstd = None self.davg = None self.place_holders = {} avg_zero = np.zeros([self.ntypes, self.ndescrpt]).astype(global_np_float_precision) std_ones = np.ones([self.ntypes, self.ndescrpt]).astype(global_np_float_precision) sub_graph = tf.Graph() with sub_graph.as_default(): name_pfx = 'd_sea_' for ii in ['coord', 'box']: self.place_holders[ii] = tf.placeholder( global_np_float_precision, [None, None], name=name_pfx + 't_' + ii) self.place_holders['type'] = tf.placeholder(tf.int32, [None, None], name=name_pfx + 't_type') self.place_holders['natoms_vec'] = tf.placeholder( tf.int32, [self.ntypes + 2], name=name_pfx + 't_natoms') self.place_holders['default_mesh'] = tf.placeholder( tf.int32, [None], name=name_pfx + 't_mesh') self.stat_descrpt, descrpt_deriv, rij, nlist \ = op_module.descrpt_se_a(self.place_holders['coord'], self.place_holders['type'], self.place_holders['natoms_vec'], self.place_holders['box'], self.place_holders['default_mesh'], tf.constant(avg_zero), tf.constant(std_ones), rcut_a = self.rcut_a, rcut_r = self.rcut_r, rcut_r_smth = self.rcut_r_smth, sel_a = self.sel_a, sel_r = self.sel_r) self.sub_sess = tf.Session(graph=sub_graph, config=default_tf_session_config)
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.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 = '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)