diff --git a/deepmd/descriptor/se_a.py b/deepmd/descriptor/se_a.py index b568d8be71..9b54daaa94 100644 --- a/deepmd/descriptor/se_a.py +++ b/deepmd/descriptor/se_a.py @@ -164,10 +164,7 @@ def __init__ (self, self.embedding_net_variables = None self.mixed_prec = None self.place_holders = {} - nei_type = np.array([]) - for ii in range(self.ntypes): - nei_type = np.append(nei_type, ii * np.ones(self.sel_a[ii])) # like a mask - self.nei_type = tf.constant(nei_type, dtype = tf.int32) + self.nei_type = np.repeat(np.arange(self.ntypes), self.sel_a) # like a mask 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) @@ -673,8 +670,9 @@ def _concat_type_embedding( embedding: environment of each atom represented by embedding. ''' - te_out_dim = type_embedding.get_shape().as_list()[-1] - nei_embed = tf.nn.embedding_lookup(type_embedding,tf.cast(self.nei_type,dtype=tf.int32)) # shape is [self.nnei, 1+te_out_dim] + te_out_dim = type_embedding.get_shape().as_list()[-1] + self.t_nei_type = tf.constant(self.nei_type, dtype=tf.int32) + nei_embed = tf.nn.embedding_lookup(type_embedding,tf.cast(self.t_nei_type,dtype=tf.int32)) # shape is [self.nnei, 1+te_out_dim] nei_embed = tf.tile(nei_embed,(nframes*natoms[0],1)) # shape is [nframes*natoms[0]*self.nnei, te_out_dim] nei_embed = tf.reshape(nei_embed,[-1,te_out_dim]) embedding_input = tf.concat([xyz_scatter,nei_embed],1) # shape is [nframes*natoms[0]*self.nnei, 1+te_out_dim] diff --git a/deepmd/descriptor/se_atten.py b/deepmd/descriptor/se_atten.py index ead4f93d46..ee2ab39a7f 100644 --- a/deepmd/descriptor/se_atten.py +++ b/deepmd/descriptor/se_atten.py @@ -15,6 +15,7 @@ from deepmd.utils.type_embed import embed_atom_type from deepmd.utils.sess import run_sess from deepmd.utils.graph import load_graph_def, get_tensor_by_name_from_graph, get_tensor_by_name +from deepmd.utils.graph import get_attention_layer_variables_from_graph_def from deepmd.utils.errors import GraphWithoutTensorError from .descriptor import Descriptor from .se_a import DescrptSeA @@ -117,6 +118,9 @@ def __init__(self, self.sel_all_r = [0] 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) + self.beta = np.zeros([self.attn_layer, self.filter_neuron[-1]]).astype(GLOBAL_NP_FLOAT_PRECISION) + self.gamma = np.ones([self.attn_layer, self.filter_neuron[-1]]).astype(GLOBAL_NP_FLOAT_PRECISION) + self.attention_layer_variables = None sub_graph = tf.Graph() with sub_graph.as_default(): name_pfx = 'd_sea_' @@ -305,10 +309,6 @@ def build(self, self.attn_weight = [None for i in range(self.attn_layer)] self.angular_weight = [None for i in range(self.attn_layer)] self.attn_weight_final = [None for i in range(self.attn_layer)] - self.G = None - self.qs = [None for i in range(self.attn_layer)] - self.ks = [None for i in range(self.attn_layer)] - self.vs = [None for i in range(self.attn_layer)] self.descrpt, self.descrpt_deriv, self.rij, self.nlist, self.nei_type_vec, self.nmask \ = op_module.prod_env_mat_a_mix(coord, @@ -365,8 +365,8 @@ def _pass_filter(self, inputs_i = inputs inputs_i = tf.reshape(inputs_i, [-1, self.ndescrpt]) type_i = -1 - layer, qmat = self._filter(inputs_i, type_i, natoms, name='filter_type_all' + suffix, reuse=reuse, - trainable=trainable, activation_fn=self.filter_activation_fn, + layer, qmat = self._filter(inputs_i, type_i, natoms, name='filter_type_all' + suffix, suffix=suffix, + reuse=reuse, trainable=trainable, activation_fn=self.filter_activation_fn, type_embedding=type_embedding, atype=atype) layer = tf.reshape(layer, [tf.shape(inputs)[0], natoms[0], self.get_dim_out()]) qmat = tf.reshape(qmat, [tf.shape(inputs)[0], natoms[0], self.get_dim_rot_mat_1() * 3]) @@ -508,7 +508,8 @@ def _feedforward(self, input_xyz, d_in, d_mid): activation_fn=None, precision=self.filter_precision, trainable=True, - uniform_seed=self.uniform_seed)) + uniform_seed=self.uniform_seed, + initial_variables=self.attention_layer_variables)) input_xyz = one_layer( input_xyz, d_in, @@ -518,7 +519,8 @@ def _feedforward(self, input_xyz, d_in, d_mid): activation_fn=None, precision=self.filter_precision, trainable=True, - uniform_seed=self.uniform_seed) + uniform_seed=self.uniform_seed, + initial_variables=self.attention_layer_variables) input_xyz += residual input_xyz = tf.keras.layers.LayerNormalization()(input_xyz) return input_xyz @@ -553,75 +555,75 @@ def _attention_layers( input_r, dotr=False, do_mask=False, - trainable=True + trainable=True, + suffix='' ): sd_k = tf.sqrt(tf.cast(1., dtype=self.filter_precision)) - self.G = tf.reshape(input_xyz, (-1, shape_i[1] // 4, outputs_size[-1]))[0] for i in range(layer_num): - with tf.variable_scope('attention_layer{}_'.format(i), reuse=tf.AUTO_REUSE): + name = 'attention_layer_{}{}'.format(i, suffix) + with tf.variable_scope(name, reuse=tf.AUTO_REUSE): # input_xyz_in = tf.nn.l2_normalize(input_xyz, -1) Q_c = one_layer( input_xyz, self.att_n, name='c_query', + scope=name+'/', reuse=tf.AUTO_REUSE, seed=self.seed, activation_fn=None, precision=self.filter_precision, trainable=trainable, - uniform_seed=self.uniform_seed) + uniform_seed=self.uniform_seed, + initial_variables=self.attention_layer_variables) K_c = one_layer( input_xyz, self.att_n, name='c_key', + scope=name+'/', reuse=tf.AUTO_REUSE, seed=self.seed, activation_fn=None, precision=self.filter_precision, trainable=trainable, - uniform_seed=self.uniform_seed) + uniform_seed=self.uniform_seed, + initial_variables=self.attention_layer_variables) V_c = one_layer( input_xyz, self.att_n, name='c_value', + scope=name+'/', reuse=tf.AUTO_REUSE, seed=self.seed, activation_fn=None, precision=self.filter_precision, trainable=trainable, - uniform_seed=self.uniform_seed) + uniform_seed=self.uniform_seed, + initial_variables=self.attention_layer_variables) # # natom x nei_type_i x out_size # xyz_scatter = tf.reshape(xyz_scatter, (-1, shape_i[1] // 4, outputs_size[-1])) # natom x nei_type_i x att_n Q_c = tf.nn.l2_normalize(tf.reshape(Q_c, (-1, shape_i[1] // 4, self.att_n)), -1) K_c = tf.nn.l2_normalize(tf.reshape(K_c, (-1, shape_i[1] // 4, self.att_n)), -1) V_c = tf.nn.l2_normalize(tf.reshape(V_c, (-1, shape_i[1] // 4, self.att_n)), -1) - # Q_c = tf.reshape(Q_c, (-1, shape_i[1] // 4, self.att_n)) - # K_c = tf.reshape(K_c, (-1, shape_i[1] // 4, self.att_n)) - # V_c = tf.reshape(V_c, (-1, shape_i[1] // 4, self.att_n)) - self.qs[i] = Q_c[0] - self.ks[i] = K_c[0] - self.vs[i] = V_c[0] input_att = self._scaled_dot_attn(Q_c, K_c, V_c, sd_k, input_r, dotr=dotr, do_mask=do_mask, layer=i) input_att = tf.reshape(input_att, (-1, self.att_n)) - # A_c = tf.nn.softmax(tf.matmul(Q_c, K_c, transpose_b=True)/sd_k) - # # (natom x nei_type_i) x att_n - # input_att = tf.reshape(tf.matmul(A_c, V_c), (-1, self.att_n)) - # (natom x nei_type_i) x out_size input_xyz += one_layer( input_att, outputs_size[-1], name='c_out', + scope=name+'/', reuse=tf.AUTO_REUSE, seed=self.seed, activation_fn=None, precision=self.filter_precision, trainable=trainable, - uniform_seed=self.uniform_seed) - input_xyz = tf.keras.layers.LayerNormalization()(input_xyz) + uniform_seed=self.uniform_seed, + initial_variables=self.attention_layer_variables) + input_xyz = tf.keras.layers.LayerNormalization(beta_initializer=tf.constant_initializer(self.beta[i]), + gamma_initializer=tf.constant_initializer(self.gamma[i]))(input_xyz) # input_xyz = self._feedforward(input_xyz, outputs_size[-1], self.att_n) return input_xyz @@ -688,7 +690,7 @@ def _filter_lower( # natom x nei_type_i x out_size xyz_scatter_att = tf.reshape( self._attention_layers(xyz_scatter, self.attn_layer, shape_i, outputs_size, input_r, - dotr=self.attn_dotr, do_mask=self.attn_mask, trainable=trainable), + dotr=self.attn_dotr, do_mask=self.attn_mask, trainable=trainable, suffix=suffix), (-1, shape_i[1] // 4, outputs_size[-1])) # xyz_scatter = tf.reshape(xyz_scatter, (-1, shape_i[1] // 4, outputs_size[-1])) else: @@ -712,6 +714,7 @@ def _filter( activation_fn=tf.nn.tanh, stddev=1.0, bavg=0.0, + suffix='', name='linear', reuse=None, trainable=True): @@ -745,6 +748,7 @@ def _filter( stddev=stddev, bavg=bavg, trainable=trainable, + suffix=suffix, name=name, reuse=reuse, atype=atype) @@ -775,3 +779,31 @@ def _filter( result = tf.reshape(result, [-1, outputs_size_2 * outputs_size[-1]]) return result, qmat + + def init_variables(self, + graph: tf.Graph, + graph_def: tf.GraphDef, + suffix: str = "", + ) -> None: + """ + Init the embedding net variables with the given dict + + Parameters + ---------- + graph : tf.Graph + The input frozen model graph + graph_def : tf.GraphDef + The input frozen model graph_def + suffix : str, optional + The suffix of the scope + """ + super().init_variables(graph=graph, graph_def=graph_def, suffix=suffix) + self.attention_layer_variables = get_attention_layer_variables_from_graph_def(graph_def, suffix=suffix) + if self.attn_layer > 0: + self.beta[0] = self.attention_layer_variables['attention_layer_0{}/layer_normalization/beta'.format(suffix)] + self.gamma[0] = self.attention_layer_variables['attention_layer_0{}/layer_normalization/gamma'.format(suffix)] + for i in range(1, self.attn_layer): + self.beta[i] = self.attention_layer_variables[ + 'attention_layer_{}{}/layer_normalization_{}/beta'.format(i, suffix, i)] + self.gamma[i] = self.attention_layer_variables[ + 'attention_layer_{}{}/layer_normalization_{}/gamma'.format(i, suffix, i)] diff --git a/deepmd/env.py b/deepmd/env.py index f0d6e1da4e..36c4b4b7c2 100644 --- a/deepmd/env.py +++ b/deepmd/env.py @@ -44,6 +44,8 @@ "TRANSFER_PATTERN", "FITTING_NET_PATTERN", "EMBEDDING_NET_PATTERN", + "TYPE_EMBEDDING_PATTERN", + "ATTENTION_LAYER_PATTERN", "TF_VERSION" ] @@ -59,18 +61,26 @@ r"filter_type_\d+/matrix_\d+_\d+|" r"filter_type_\d+/bias_\d+_\d+|" r"filter_type_\d+/idt_\d+_\d+|" + r"filter_type_all/matrix_\d+|" r"filter_type_all/matrix_\d+_\d+|" r"filter_type_all/matrix_\d+_\d+_\d+|" + r"filter_type_all/bias_\d+|" r"filter_type_all/bias_\d+_\d+|" r"filter_type_all/bias_\d+_\d+_\d+|" + r"filter_type_all/idt_\d+|" r"filter_type_all/idt_\d+_\d+|" ) FITTING_NET_PATTERN = str( + r"layer_\d+/matrix|" r"layer_\d+_type_\d+/matrix|" + r"layer_\d+/bias|" r"layer_\d+_type_\d+/bias|" + r"layer_\d+/idt|" r"layer_\d+_type_\d+/idt|" + r"final_layer/matrix|" r"final_layer_type_\d+/matrix|" + r"final_layer/bias|" r"final_layer_type_\d+/bias|" ) @@ -80,6 +90,21 @@ r"type_embed_net+/idt_\d+|" ) +ATTENTION_LAYER_PATTERN = str( + r"attention_layer_\d+/c_query/matrix|" + r"attention_layer_\d+/c_query/bias|" + r"attention_layer_\d+/c_key/matrix|" + r"attention_layer_\d+/c_key/bias|" + r"attention_layer_\d+/c_value/matrix|" + r"attention_layer_\d+/c_value/bias|" + r"attention_layer_\d+/c_out/matrix|" + r"attention_layer_\d+/c_out/bias|" + r"attention_layer_\d+/layer_normalization/beta|" + r"attention_layer_\d+/layer_normalization/gamma|" + r"attention_layer_\d+/layer_normalization_\d+/beta|" + r"attention_layer_\d+/layer_normalization_\d+/gamma|" +) + TRANSFER_PATTERN = \ EMBEDDING_NET_PATTERN + \ FITTING_NET_PATTERN + \ diff --git a/deepmd/fit/ener.py b/deepmd/fit/ener.py index 64258d88ec..47e9589cf3 100644 --- a/deepmd/fit/ener.py +++ b/deepmd/fit/ener.py @@ -9,6 +9,7 @@ from deepmd.utils.network import one_layer as one_layer_deepmd from deepmd.utils.type_embed import embed_atom_type from deepmd.utils.graph import get_fitting_net_variables_from_graph_def, load_graph_def, get_tensor_by_name_from_graph +from deepmd.utils.errors import GraphWithoutTensorError from deepmd.fit.fitting import Fitting from deepmd.env import global_cvt_2_tf_float @@ -400,6 +401,8 @@ def build (self, if input_dict is None: input_dict = {} bias_atom_e = self.bias_atom_e + type_embedding = input_dict.get('type_embedding', None) + atype = input_dict.get('atype', None) if self.numb_fparam > 0: if self.fparam_avg is None: self.fparam_avg = 0. @@ -418,9 +421,10 @@ def build (self, t_daparam = tf.constant(self.numb_aparam, name = 'daparam', dtype = tf.int32) - self.t_bias_atom_e = tf.get_variable('t_bias_atom_e', + if type_embedding is not None: + self.t_bias_atom_e = tf.get_variable('t_bias_atom_e', self.bias_atom_e.shape, - dtype=GLOBAL_TF_FLOAT_PRECISION, + dtype=self.fitting_precision, trainable=False, initializer=tf.constant_initializer(self.bias_atom_e)) if self.numb_fparam > 0: @@ -471,9 +475,7 @@ def build (self, aparam = tf.reshape(aparam, [-1, self.numb_aparam]) aparam = (aparam - t_aparam_avg) * t_aparam_istd aparam = tf.reshape(aparam, [-1, self.numb_aparam * natoms[0]]) - - type_embedding = input_dict.get('type_embedding', None) - atype = input_dict.get('atype', None) + if type_embedding is not None: atype_nall = tf.reshape(atype, [-1, natoms[1]]) self.atype_nloc = tf.reshape(tf.slice(atype_nall, [0, 0], [-1, natoms[0]]), [-1]) ## lammps will make error @@ -570,6 +572,11 @@ def init_variables(self, if self.numb_aparam > 0: self.aparam_avg = get_tensor_by_name_from_graph(graph, 'fitting_attr%s/t_aparam_avg' % suffix) self.aparam_inv_std = get_tensor_by_name_from_graph(graph, 'fitting_attr%s/t_aparam_istd' % suffix) + try: + self.bias_atom_e = get_tensor_by_name_from_graph(graph, 'fitting_attr%s/t_bias_atom_e' % suffix) + except GraphWithoutTensorError: + # model without type_embedding has no t_bias_atom_e + pass def enable_compression(self, model_file: str, diff --git a/deepmd/train/trainer.py b/deepmd/train/trainer.py index 52cee2427b..79038709c1 100644 --- a/deepmd/train/trainer.py +++ b/deepmd/train/trainer.py @@ -288,7 +288,8 @@ def _init_param(self, jdata): def build (self, data = None, - stop_batch = 0) : + stop_batch = 0, + suffix = "") : self.ntypes = self.model.get_ntypes() self.stop_batch = stop_batch @@ -348,7 +349,7 @@ def build (self, self.fitting.enable_mixed_precision(self.mixed_prec) self._build_lr() - self._build_network(data) + self._build_network(data, suffix) self._build_training() @@ -358,7 +359,7 @@ def _build_lr(self): self.learning_rate = self.lr.build(self.global_step, self.stop_batch) log.info("built lr") - def _build_network(self, data): + def _build_network(self, data, suffix=""): self.place_holders = {} if self.is_compress : for kk in ['coord', 'box']: @@ -379,7 +380,7 @@ def _build_network(self, data): self.place_holders['default_mesh'], self.place_holders, self.frz_model, - suffix = "", + suffix = suffix, reuse = False) self.l2_l, self.l2_more\ diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 1f84b7fe60..ce7710d116 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -276,9 +276,9 @@ def descrpt_se_atten_args(): Argument("seed", [int, None], optional=True, doc=doc_seed), Argument("exclude_types", list, optional=True, default=[], doc=doc_exclude_types), Argument("set_davg_zero", bool, optional=True, default=False, doc=doc_set_davg_zero), - Argument("attn", int, optional=True, default=100, doc=doc_attn), - Argument("attn_layer", int, optional=True, default=4, doc=doc_attn_layer), - Argument("attn_dotr", bool, optional=True, default=False, doc=doc_attn_dotr), + Argument("attn", int, optional=True, default=128, doc=doc_attn), + Argument("attn_layer", int, optional=True, default=2, doc=doc_attn_layer), + Argument("attn_dotr", bool, optional=True, default=True, doc=doc_attn_dotr), Argument("attn_mask", bool, optional=True, default=False, doc=doc_attn_mask) ] diff --git a/deepmd/utils/graph.py b/deepmd/utils/graph.py index 10e13730c5..a8d95ebb25 100644 --- a/deepmd/utils/graph.py +++ b/deepmd/utils/graph.py @@ -1,7 +1,7 @@ import re import numpy as np from typing import Tuple, Dict -from deepmd.env import tf, EMBEDDING_NET_PATTERN, FITTING_NET_PATTERN, TYPE_EMBEDDING_PATTERN +from deepmd.env import tf, EMBEDDING_NET_PATTERN, FITTING_NET_PATTERN, TYPE_EMBEDDING_PATTERN, ATTENTION_LAYER_PATTERN from deepmd.utils.sess import run_sess from deepmd.utils.errors import GraphWithoutTensorError @@ -391,3 +391,63 @@ def get_type_embedding_net_variables_from_graph_def(graph_def: tf.GraphDef, suff tensor_value = get_tensor_by_type(node, dtype) type_embedding_net_variables[item] = np.reshape(tensor_value, tensor_shape) return type_embedding_net_variables + + +def get_attention_layer_nodes_from_graph_def(graph_def: tf.GraphDef, suffix: str = "") -> Dict: + """ + Get the attention layer nodes with the given tf.GraphDef object + + Parameters + ---------- + graph_def + The input tf.GraphDef object + suffix : str, optional + The scope suffix + + Returns + ---------- + Dict + The attention layer nodes within the given tf.GraphDef object + """ + if suffix != "": + attention_layer_pattern = ATTENTION_LAYER_PATTERN \ + .replace('/c_query', suffix + '/c_query') \ + .replace('/c_key', suffix + '/c_key') \ + .replace('/c_value', suffix + '/c_value') \ + .replace('/c_out', suffix + '/c_out') \ + .replace('/layer_normalization', suffix + '/layer_normalization') + else: + attention_layer_pattern = ATTENTION_LAYER_PATTERN + + attention_layer_nodes = get_pattern_nodes_from_graph_def(graph_def, attention_layer_pattern) + return attention_layer_nodes + + +def get_attention_layer_variables_from_graph_def(graph_def: tf.GraphDef, suffix: str = "") -> Dict: + """ + Get the attention layer variables with the given tf.GraphDef object + + Parameters + ---------- + graph_def : tf.GraphDef + The input tf.GraphDef object + suffix : str, optional + The suffix of the scope + + Returns + ---------- + Dict + The attention layer variables within the given tf.GraphDef object + """ + attention_layer_variables = {} + attention_layer_net_nodes = get_attention_layer_nodes_from_graph_def(graph_def, suffix=suffix) + for item in attention_layer_net_nodes: + node = attention_layer_net_nodes[item] + dtype = tf.as_dtype(node.dtype).as_numpy_dtype + tensor_shape = tf.TensorShape(node.tensor_shape).as_list() + if (len(tensor_shape) != 1) or (tensor_shape[0] != 1): + tensor_value = np.frombuffer(node.tensor_content, dtype=tf.as_dtype(node.dtype).as_numpy_dtype) + else: + tensor_value = get_tensor_by_type(node, dtype) + attention_layer_variables[item] = np.reshape(tensor_value, tensor_shape) + return attention_layer_variables diff --git a/deepmd/utils/network.py b/deepmd/utils/network.py index 22542adfba..47d9d83351 100644 --- a/deepmd/utils/network.py +++ b/deepmd/utils/network.py @@ -13,7 +13,8 @@ def one_layer(inputs, precision = GLOBAL_TF_FLOAT_PRECISION, stddev=1.0, bavg=0.0, - name='linear', + name='linear', + scope='', reuse=None, seed=None, use_timestep = False, @@ -36,8 +37,8 @@ def one_layer(inputs, mean=bavg, seed=seed if (seed is None or uniform_seed) else seed + 1) if initial_variables is not None: - w_initializer = tf.constant_initializer(initial_variables[name + '/matrix']) - b_initializer = tf.constant_initializer(initial_variables[name + '/bias']) + w_initializer = tf.constant_initializer(initial_variables[scope + name + '/matrix']) + b_initializer = tf.constant_initializer(initial_variables[scope + name + '/bias']) w = tf.get_variable('matrix', [shape[1], outputs_size], precision, @@ -63,7 +64,7 @@ def one_layer(inputs, mean=0.1, seed=seed if (seed is None or uniform_seed) else seed + 2) if initial_variables is not None: - idt_initializer = tf.constant_initializer(initial_variables[name + '/idt']) + idt_initializer = tf.constant_initializer(initial_variables[scope + name + '/idt']) idt = tf.get_variable('idt', [outputs_size], precision, diff --git a/source/tests/init_frz_model/data/set.000/box.npy b/source/tests/init_frz_model/data/set.000/box.npy new file mode 100644 index 0000000000..aa092a9aad Binary files /dev/null and b/source/tests/init_frz_model/data/set.000/box.npy differ diff --git a/source/tests/init_frz_model/data/set.000/coord.npy b/source/tests/init_frz_model/data/set.000/coord.npy new file mode 100644 index 0000000000..2205bfd452 Binary files /dev/null and b/source/tests/init_frz_model/data/set.000/coord.npy differ diff --git a/source/tests/init_frz_model/data/set.000/energy.npy b/source/tests/init_frz_model/data/set.000/energy.npy new file mode 100644 index 0000000000..6c1e5ce145 Binary files /dev/null and b/source/tests/init_frz_model/data/set.000/energy.npy differ diff --git a/source/tests/init_frz_model/data/set.000/force.npy b/source/tests/init_frz_model/data/set.000/force.npy new file mode 100644 index 0000000000..10698ebb59 Binary files /dev/null and b/source/tests/init_frz_model/data/set.000/force.npy differ diff --git a/source/tests/init_frz_model/data/type.raw b/source/tests/init_frz_model/data/type.raw new file mode 100644 index 0000000000..e329bb5191 --- /dev/null +++ b/source/tests/init_frz_model/data/type.raw @@ -0,0 +1,6 @@ +0 +1 +1 +0 +1 +1 \ No newline at end of file diff --git a/source/tests/init_frz_model/data/type_map.raw b/source/tests/init_frz_model/data/type_map.raw new file mode 100644 index 0000000000..e900768b1d --- /dev/null +++ b/source/tests/init_frz_model/data/type_map.raw @@ -0,0 +1,2 @@ +O +H diff --git a/source/tests/init_frz_model/input.json b/source/tests/init_frz_model/input.json new file mode 100644 index 0000000000..4c4e93fb05 --- /dev/null +++ b/source/tests/init_frz_model/input.json @@ -0,0 +1,66 @@ +{ + "_comment": " model parameters", + "model": { +"type_map": ["O", "H"], +"descriptor" :{ + "type": "se_e2_a", + "sel": [46, 92], + "rcut_smth": 0.50, + "rcut": 6.00, + "neuron": [4, 8, 16], + "resnet_dt": false, + "axis_neuron": 16, + "seed": 1, + "_comment": " that's all" +}, +"fitting_net" : { + "neuron": [20, 20, 20], + "resnet_dt": true, + "seed": 1, + "_comment": " that's all" +}, +"_comment": " that's all" + }, + + "learning_rate" :{ +"type": "exp", +"decay_steps": 5000, +"start_lr": 0.001, +"stop_lr": 3.51e-8, +"_comment": "that's all" + }, + + "loss" :{ +"type": "ener", +"start_pref_e": 0.02, +"limit_pref_e": 1, +"start_pref_f": 1000, +"limit_pref_f": 1, +"start_pref_v": 1, +"limit_pref_v": 1, +"_comment": " that's all" + }, + + "training" : { +"training_data": { + "systems": ["init_frz_model/data"], + "batch_size": "auto", + "_comment": "that's all" +}, +"validation_data":{ + "systems": ["init_frz_model/data"], + "batch_size": 1, + "numb_btch": 3, + "_comment": "that's all" +}, +"numb_steps": 1, +"seed": 10, +"disp_file": "lcurve.out", +"disp_freq": 1, +"save_freq": 1, +"_comment": "that's all" + }, + + "_comment": "that's all" +} + diff --git a/source/tests/test_init_frz_model_se_a.py b/source/tests/test_init_frz_model_se_a.py new file mode 100644 index 0000000000..cad93dbeca --- /dev/null +++ b/source/tests/test_init_frz_model_se_a.py @@ -0,0 +1,161 @@ +import os, sys, platform, shutil, dpdata, json +import numpy as np +import unittest +import subprocess as sp +from common import j_loader, tests_path +from deepmd.train.trainer import DPTrainer +from deepmd.train.run_options import RunOptions +from deepmd.utils.argcheck import normalize +from deepmd.utils.compat import update_deepmd_input +from deepmd.utils.data_system import DeepmdDataSystem + +from deepmd.env import GLOBAL_NP_FLOAT_PRECISION, tf + +if GLOBAL_NP_FLOAT_PRECISION == np.float32: + default_places = 4 +else: + default_places = 10 + + +def _file_delete(file): + if os.path.isdir(file): + os.rmdir(file) + elif os.path.isfile(file): + os.remove(file) + + +def _subprocess_run(command): + popen = sp.Popen(command.split(), shell=False, stdout=sp.PIPE, stderr=sp.STDOUT) + for line in iter(popen.stdout.readline, b''): + if hasattr(line, 'decode'): + line = line.decode('utf-8') + line = line.rstrip() + print(line) + popen.wait() + return popen.returncode + + +def _init_models(): + data_file = str(tests_path / os.path.join("init_frz_model", "data")) + frozen_model = str(tests_path / "init_frz_se_a.pb") + ckpt = str(tests_path / "init_frz_se_a.ckpt") + run_opt_ckpt = RunOptions(init_model=ckpt, log_level=20) + run_opt_frz = RunOptions(init_frz_model=frozen_model, log_level=20) + INPUT = str(tests_path / "input.json") + jdata = j_loader(str(tests_path / os.path.join("init_frz_model", "input.json"))) + jdata["training"]["training_data"]["systems"] = data_file + jdata["training"]["validation_data"]["systems"] = data_file + jdata["training"]["save_ckpt"] = ckpt + with open(INPUT, "w") as fp: + json.dump(jdata, fp, indent=4) + ret = _subprocess_run("dp train " + INPUT) + np.testing.assert_equal(ret, 0, 'DP train failed!') + ret = _subprocess_run("dp freeze -c " + str(tests_path) + " -o " + frozen_model) + np.testing.assert_equal(ret, 0, 'DP freeze failed!') + + jdata = update_deepmd_input(jdata, warning=True, dump="input_v2_compat.json") + jdata = normalize(jdata) + model_ckpt = DPTrainer(jdata, run_opt=run_opt_ckpt) + model_frz = DPTrainer(jdata, run_opt=run_opt_frz) + rcut = model_ckpt.model.get_rcut() + type_map = model_ckpt.model.get_type_map() + data = DeepmdDataSystem( + systems=[data_file], + batch_size=1, + test_size=1, + rcut=rcut, + type_map=type_map, + trn_all_set=True + ) + data_requirement = {'energy': {'ndof': 1, + 'atomic': False, + 'must': False, + 'high_prec': True, + 'type_sel': None, + 'repeat': 1, + 'default': 0.0}, + 'force': {'ndof': 3, + 'atomic': True, + 'must': False, + 'high_prec': False, + 'type_sel': None, + 'repeat': 1, + 'default': 0.0}, + 'virial': {'ndof': 9, + 'atomic': False, + 'must': False, + 'high_prec': False, + 'type_sel': None, + 'repeat': 1, + 'default': 0.0}, + 'atom_ener': {'ndof': 1, + 'atomic': True, + 'must': False, + 'high_prec': False, + 'type_sel': None, + 'repeat': 1, + 'default': 0.0}, + 'atom_pref': {'ndof': 1, + 'atomic': True, + 'must': False, + 'high_prec': False, + 'type_sel': None, + 'repeat': 3, + 'default': 0.0}} + data.add_dict(data_requirement) + stop_batch = jdata["training"]["numb_steps"] + + return INPUT, ckpt, frozen_model, model_ckpt, model_frz, data, stop_batch + + +INPUT, CKPT, FROZEN_MODEL, CKPT_TRAINER, FRZ_TRAINER, VALID_DATA, STOP_BATCH = _init_models() + + +class TestInitFrzModelA(unittest.TestCase): + @classmethod + def setUpClass(self): + self.dp_ckpt = CKPT_TRAINER + self.dp_frz = FRZ_TRAINER + self.valid_data = VALID_DATA + self.stop_batch = STOP_BATCH + + @classmethod + def tearDownClass(self): + _file_delete(INPUT) + _file_delete(FROZEN_MODEL) + _file_delete("out.json") + _file_delete(str(tests_path / "checkpoint")) + _file_delete(CKPT+".meta") + _file_delete(CKPT+".index") + _file_delete(CKPT+".data-00000-of-00001") + _file_delete(CKPT+"-0.meta") + _file_delete(CKPT+"-0.index") + _file_delete(CKPT+"-0.data-00000-of-00001") + _file_delete(CKPT+"-1.meta") + _file_delete(CKPT+"-1.index") + _file_delete(CKPT+"-1.data-00000-of-00001") + _file_delete("input_v2_compat.json") + _file_delete("lcurve.out") + + def test_single_frame(self): + valid_batch = self.valid_data.get_batch() + natoms = valid_batch['natoms_vec'] + tf.reset_default_graph() + self.dp_ckpt.build(self.valid_data, self.stop_batch) + self.dp_ckpt._init_session() + feed_dict_ckpt = self.dp_ckpt.get_feed_dict(valid_batch, is_training=False) + ckpt_rmse_ckpt = self.dp_ckpt.loss.eval(self.dp_ckpt.sess, feed_dict_ckpt, natoms) + tf.reset_default_graph() + + self.dp_frz.build(self.valid_data, self.stop_batch) + self.dp_frz._init_session() + feed_dict_frz = self.dp_frz.get_feed_dict(valid_batch, is_training=False) + ckpt_rmse_frz = self.dp_frz.loss.eval(self.dp_frz.sess, feed_dict_frz, natoms) + tf.reset_default_graph() + + # check values + np.testing.assert_almost_equal(ckpt_rmse_ckpt['rmse_e'], ckpt_rmse_frz['rmse_e'], default_places) + np.testing.assert_almost_equal(ckpt_rmse_ckpt['rmse_f'], ckpt_rmse_frz['rmse_f'], default_places) + np.testing.assert_almost_equal(ckpt_rmse_ckpt['rmse_v'], ckpt_rmse_frz['rmse_v'], default_places) + + diff --git a/source/tests/test_init_frz_model_se_a_type.py b/source/tests/test_init_frz_model_se_a_type.py new file mode 100644 index 0000000000..2e7a693e3f --- /dev/null +++ b/source/tests/test_init_frz_model_se_a_type.py @@ -0,0 +1,165 @@ +import os, sys, platform, shutil, dpdata, json +import numpy as np +import unittest +import subprocess as sp +from common import j_loader, tests_path +from deepmd.train.trainer import DPTrainer +from deepmd.train.run_options import RunOptions +from deepmd.utils.argcheck import normalize +from deepmd.utils.compat import update_deepmd_input +from deepmd.utils.data_system import DeepmdDataSystem + +from deepmd.env import GLOBAL_NP_FLOAT_PRECISION, tf + +if GLOBAL_NP_FLOAT_PRECISION == np.float32: + default_places = 4 +else: + default_places = 10 + + +def _file_delete(file): + if os.path.isdir(file): + os.rmdir(file) + elif os.path.isfile(file): + os.remove(file) + + +def _subprocess_run(command): + popen = sp.Popen(command.split(), shell=False, stdout=sp.PIPE, stderr=sp.STDOUT) + for line in iter(popen.stdout.readline, b''): + if hasattr(line, 'decode'): + line = line.decode('utf-8') + line = line.rstrip() + print(line) + popen.wait() + return popen.returncode + + +def _init_models(): + data_file = str(tests_path / os.path.join("init_frz_model", "data")) + frozen_model = str(tests_path / "init_frz_se_a_type.pb") + ckpt = str(tests_path / "init_frz_se_a_type.ckpt") + run_opt_ckpt = RunOptions(init_model=ckpt, log_level=20) + run_opt_frz = RunOptions(init_frz_model=frozen_model, log_level=20) + INPUT = str(tests_path / "input.json") + jdata = j_loader(str(tests_path / os.path.join("init_frz_model", "input.json"))) + jdata["training"]["training_data"]["systems"] = data_file + jdata["training"]["validation_data"]["systems"] = data_file + jdata["training"]["save_ckpt"] = ckpt + type_embed = {} + type_embed['neuron'] = [2, 4, 8] + type_embed['resnet_dt'] = False + jdata['model']["type_embedding"] = type_embed + with open(INPUT, "w") as fp: + json.dump(jdata, fp, indent=4) + ret = _subprocess_run("dp train " + INPUT) + np.testing.assert_equal(ret, 0, 'DP train failed!') + ret = _subprocess_run("dp freeze -c " + str(tests_path) + " -o " + frozen_model) + np.testing.assert_equal(ret, 0, 'DP freeze failed!') + + jdata = update_deepmd_input(jdata, warning=True, dump="input_v2_compat.json") + jdata = normalize(jdata) + model_ckpt = DPTrainer(jdata, run_opt=run_opt_ckpt) + model_frz = DPTrainer(jdata, run_opt=run_opt_frz) + rcut = model_ckpt.model.get_rcut() + type_map = model_ckpt.model.get_type_map() + data = DeepmdDataSystem( + systems=[data_file], + batch_size=1, + test_size=1, + rcut=rcut, + type_map=type_map, + trn_all_set=True + ) + data_requirement = {'energy': {'ndof': 1, + 'atomic': False, + 'must': False, + 'high_prec': True, + 'type_sel': None, + 'repeat': 1, + 'default': 0.0}, + 'force': {'ndof': 3, + 'atomic': True, + 'must': False, + 'high_prec': False, + 'type_sel': None, + 'repeat': 1, + 'default': 0.0}, + 'virial': {'ndof': 9, + 'atomic': False, + 'must': False, + 'high_prec': False, + 'type_sel': None, + 'repeat': 1, + 'default': 0.0}, + 'atom_ener': {'ndof': 1, + 'atomic': True, + 'must': False, + 'high_prec': False, + 'type_sel': None, + 'repeat': 1, + 'default': 0.0}, + 'atom_pref': {'ndof': 1, + 'atomic': True, + 'must': False, + 'high_prec': False, + 'type_sel': None, + 'repeat': 3, + 'default': 0.0}} + data.add_dict(data_requirement) + stop_batch = jdata["training"]["numb_steps"] + + return INPUT, ckpt, frozen_model, model_ckpt, model_frz, data, stop_batch + + +INPUT, CKPT, FROZEN_MODEL, CKPT_TRAINER, FRZ_TRAINER, VALID_DATA, STOP_BATCH = _init_models() + + +class TestInitFrzModelAType(unittest.TestCase): + @classmethod + def setUpClass(self): + self.dp_ckpt = CKPT_TRAINER + self.dp_frz = FRZ_TRAINER + self.valid_data = VALID_DATA + self.stop_batch = STOP_BATCH + + @classmethod + def tearDownClass(self): + _file_delete(INPUT) + _file_delete(FROZEN_MODEL) + _file_delete("out.json") + _file_delete(str(tests_path / "checkpoint")) + _file_delete(CKPT+".meta") + _file_delete(CKPT+".index") + _file_delete(CKPT+".data-00000-of-00001") + _file_delete(CKPT+"-0.meta") + _file_delete(CKPT+"-0.index") + _file_delete(CKPT+"-0.data-00000-of-00001") + _file_delete(CKPT+"-1.meta") + _file_delete(CKPT+"-1.index") + _file_delete(CKPT+"-1.data-00000-of-00001") + _file_delete("input_v2_compat.json") + _file_delete("lcurve.out") + + def test_single_frame(self): + valid_batch = self.valid_data.get_batch() + natoms = valid_batch['natoms_vec'] + tf.reset_default_graph() + self.dp_ckpt.build(self.valid_data, self.stop_batch) + self.dp_ckpt._init_session() + feed_dict_ckpt = self.dp_ckpt.get_feed_dict(valid_batch, is_training=False) + ckpt_rmse_ckpt = self.dp_ckpt.loss.eval(self.dp_ckpt.sess, feed_dict_ckpt, natoms) + tf.reset_default_graph() + + self.dp_frz.build(self.valid_data, self.stop_batch) + self.dp_frz._init_session() + feed_dict_frz = self.dp_frz.get_feed_dict(valid_batch, is_training=False) + ckpt_rmse_frz = self.dp_frz.loss.eval(self.dp_frz.sess, feed_dict_frz, natoms) + tf.reset_default_graph() + + # check values + np.testing.assert_almost_equal(ckpt_rmse_ckpt['rmse_e'], ckpt_rmse_frz['rmse_e'], default_places) + np.testing.assert_almost_equal(ckpt_rmse_ckpt['rmse_f'], ckpt_rmse_frz['rmse_f'], default_places) + np.testing.assert_almost_equal(ckpt_rmse_ckpt['rmse_v'], ckpt_rmse_frz['rmse_v'], default_places) + + diff --git a/source/tests/test_init_frz_model_se_atten.py b/source/tests/test_init_frz_model_se_atten.py new file mode 100644 index 0000000000..0decb022e9 --- /dev/null +++ b/source/tests/test_init_frz_model_se_atten.py @@ -0,0 +1,167 @@ +import os, sys, platform, shutil, dpdata, json +import numpy as np +import unittest +import subprocess as sp +from common import j_loader, tests_path +from deepmd.train.trainer import DPTrainer +from deepmd.train.run_options import RunOptions +from deepmd.utils.argcheck import normalize +from deepmd.utils.compat import update_deepmd_input +from deepmd.utils.data_system import DeepmdDataSystem + +from deepmd.env import GLOBAL_NP_FLOAT_PRECISION, tf +from packaging.version import parse as parse_version + +if GLOBAL_NP_FLOAT_PRECISION == np.float32: + default_places = 4 +else: + default_places = 10 + + +def _file_delete(file): + if os.path.isdir(file): + os.rmdir(file) + elif os.path.isfile(file): + os.remove(file) + + +def _subprocess_run(command): + popen = sp.Popen(command.split(), shell=False, stdout=sp.PIPE, stderr=sp.STDOUT) + for line in iter(popen.stdout.readline, b''): + if hasattr(line, 'decode'): + line = line.decode('utf-8') + line = line.rstrip() + print(line) + popen.wait() + return popen.returncode + + +def _init_models(): + data_file = str(tests_path / os.path.join("init_frz_model", "data")) + frozen_model = str(tests_path / "init_frz_se_atten.pb") + ckpt = str(tests_path / "init_frz_se_atten.ckpt") + run_opt_ckpt = RunOptions(init_model=ckpt, log_level=20) + run_opt_frz = RunOptions(init_frz_model=frozen_model, log_level=20) + INPUT = str(tests_path / "input.json") + jdata = j_loader(str(tests_path / os.path.join("init_frz_model", "input.json"))) + jdata["training"]["training_data"]["systems"] = data_file + jdata["training"]["validation_data"]["systems"] = data_file + jdata["training"]["save_ckpt"] = ckpt + jdata['model']["descriptor"]['type'] = 'se_atten' + jdata['model']["descriptor"]['sel'] = 120 + with open(INPUT, "w") as fp: + json.dump(jdata, fp, indent=4) + ret = _subprocess_run("dp train " + INPUT) + np.testing.assert_equal(ret, 0, 'DP train failed!') + ret = _subprocess_run("dp freeze -c " + str(tests_path) + " -o " + frozen_model) + np.testing.assert_equal(ret, 0, 'DP freeze failed!') + + jdata = update_deepmd_input(jdata, warning=True, dump="input_v2_compat.json") + jdata = normalize(jdata) + model_ckpt = DPTrainer(jdata, run_opt=run_opt_ckpt) + model_frz = DPTrainer(jdata, run_opt=run_opt_frz) + rcut = model_ckpt.model.get_rcut() + type_map = model_ckpt.model.get_type_map() + data = DeepmdDataSystem( + systems=[data_file], + batch_size=1, + test_size=1, + rcut=rcut, + type_map=type_map, + trn_all_set=True + ) + data_requirement = {'energy': {'ndof': 1, + 'atomic': False, + 'must': False, + 'high_prec': True, + 'type_sel': None, + 'repeat': 1, + 'default': 0.0}, + 'force': {'ndof': 3, + 'atomic': True, + 'must': False, + 'high_prec': False, + 'type_sel': None, + 'repeat': 1, + 'default': 0.0}, + 'virial': {'ndof': 9, + 'atomic': False, + 'must': False, + 'high_prec': False, + 'type_sel': None, + 'repeat': 1, + 'default': 0.0}, + 'atom_ener': {'ndof': 1, + 'atomic': True, + 'must': False, + 'high_prec': False, + 'type_sel': None, + 'repeat': 1, + 'default': 0.0}, + 'atom_pref': {'ndof': 1, + 'atomic': True, + 'must': False, + 'high_prec': False, + 'type_sel': None, + 'repeat': 3, + 'default': 0.0}} + data.add_dict(data_requirement) + stop_batch = jdata["training"]["numb_steps"] + + return INPUT, ckpt, frozen_model, model_ckpt, model_frz, data, stop_batch + + +if not parse_version(tf.__version__) < parse_version("1.15"): + INPUT, CKPT, FROZEN_MODEL, CKPT_TRAINER, FRZ_TRAINER, VALID_DATA, STOP_BATCH = _init_models() + + +@unittest.skipIf(parse_version(tf.__version__) < parse_version("1.15"), + f"The current tf version {tf.__version__} is too low to run the new testing model.") +class TestInitFrzModelAtten(unittest.TestCase): + @classmethod + def setUpClass(self): + self.dp_ckpt = CKPT_TRAINER + self.dp_frz = FRZ_TRAINER + self.valid_data = VALID_DATA + self.stop_batch = STOP_BATCH + + @classmethod + def tearDownClass(self): + _file_delete(INPUT) + _file_delete(FROZEN_MODEL) + _file_delete("out.json") + _file_delete(str(tests_path / "checkpoint")) + _file_delete(CKPT+".meta") + _file_delete(CKPT+".index") + _file_delete(CKPT+".data-00000-of-00001") + _file_delete(CKPT+"-0.meta") + _file_delete(CKPT+"-0.index") + _file_delete(CKPT+"-0.data-00000-of-00001") + _file_delete(CKPT+"-1.meta") + _file_delete(CKPT+"-1.index") + _file_delete(CKPT+"-1.data-00000-of-00001") + _file_delete("input_v2_compat.json") + _file_delete("lcurve.out") + + def test_single_frame(self): + valid_batch = self.valid_data.get_batch() + natoms = valid_batch['natoms_vec'] + tf.reset_default_graph() + self.dp_ckpt.build(self.valid_data, self.stop_batch) + self.dp_ckpt._init_session() + feed_dict_ckpt = self.dp_ckpt.get_feed_dict(valid_batch, is_training=False) + ckpt_rmse_ckpt = self.dp_ckpt.loss.eval(self.dp_ckpt.sess, feed_dict_ckpt, natoms) + tf.reset_default_graph() + + self.dp_frz.build(self.valid_data, self.stop_batch) + self.dp_frz._init_session() + feed_dict_frz = self.dp_frz.get_feed_dict(valid_batch, is_training=False) + ckpt_rmse_frz = self.dp_frz.loss.eval(self.dp_frz.sess, feed_dict_frz, natoms) + tf.reset_default_graph() + + # check values + np.testing.assert_almost_equal(ckpt_rmse_ckpt['rmse_e'], ckpt_rmse_frz['rmse_e'], default_places) + np.testing.assert_almost_equal(ckpt_rmse_ckpt['rmse_f'], ckpt_rmse_frz['rmse_f'], default_places) + np.testing.assert_almost_equal(ckpt_rmse_ckpt['rmse_v'], ckpt_rmse_frz['rmse_v'], default_places) + + diff --git a/source/tests/test_init_frz_model_se_r.py b/source/tests/test_init_frz_model_se_r.py new file mode 100644 index 0000000000..f74a0ea222 --- /dev/null +++ b/source/tests/test_init_frz_model_se_r.py @@ -0,0 +1,167 @@ +import os, sys, platform, shutil, dpdata, json +import numpy as np +import unittest +import subprocess as sp +from common import j_loader, tests_path +from deepmd.train.trainer import DPTrainer +from deepmd.train.run_options import RunOptions +from deepmd.utils.argcheck import normalize +from deepmd.utils.compat import update_deepmd_input +from deepmd.utils.data_system import DeepmdDataSystem + +from deepmd.env import GLOBAL_NP_FLOAT_PRECISION, tf + +if GLOBAL_NP_FLOAT_PRECISION == np.float32: + default_places = 4 +else: + default_places = 10 + + +def _file_delete(file): + if os.path.isdir(file): + os.rmdir(file) + elif os.path.isfile(file): + os.remove(file) + + +def _subprocess_run(command): + popen = sp.Popen(command.split(), shell=False, stdout=sp.PIPE, stderr=sp.STDOUT) + for line in iter(popen.stdout.readline, b''): + if hasattr(line, 'decode'): + line = line.decode('utf-8') + line = line.rstrip() + print(line) + popen.wait() + return popen.returncode + + +def _init_models(): + data_file = str(tests_path / os.path.join("init_frz_model", "data")) + frozen_model = str(tests_path / "init_frz_se_r.pb") + ckpt = str(tests_path / "init_frz_se_r.ckpt") + run_opt_ckpt = RunOptions(init_model=ckpt, log_level=20) + run_opt_frz = RunOptions(init_frz_model=frozen_model, log_level=20) + INPUT = str(tests_path / "input.json") + jdata = j_loader(str(tests_path / os.path.join("init_frz_model", "input.json"))) + jdata["model"]["descriptor"] = {} + jdata["model"]["descriptor"]["type"] = "se_e2_r" + jdata["model"]["descriptor"]["sel"] = [46, 92] + jdata["model"]["descriptor"]["rcut_smth"] = 0.5 + jdata["model"]["descriptor"]["rcut"] = 6.0 + jdata["model"]["descriptor"]["neuron"] = [5, 10, 20] + jdata["model"]["descriptor"]["resnet_dt"] = False + jdata["model"]["descriptor"]["seed"] = 1 + jdata["training"]["training_data"]["systems"] = data_file + jdata["training"]["validation_data"]["systems"] = data_file + jdata["training"]["save_ckpt"] = ckpt + with open(INPUT, "w") as fp: + json.dump(jdata, fp, indent=4) + ret = _subprocess_run("dp train " + INPUT) + np.testing.assert_equal(ret, 0, 'DP train failed!') + ret = _subprocess_run("dp freeze -c " + str(tests_path) + " -o " + frozen_model) + np.testing.assert_equal(ret, 0, 'DP freeze failed!') + + jdata = update_deepmd_input(jdata, warning=True, dump="input_v2_compat.json") + jdata = normalize(jdata) + model_ckpt = DPTrainer(jdata, run_opt=run_opt_ckpt) + model_frz = DPTrainer(jdata, run_opt=run_opt_frz) + rcut = model_ckpt.model.get_rcut() + type_map = model_ckpt.model.get_type_map() + data = DeepmdDataSystem( + systems=[data_file], + batch_size=1, + test_size=1, + rcut=rcut, + type_map=type_map, + trn_all_set=True + ) + data_requirement = {'energy': {'ndof': 1, + 'atomic': False, + 'must': False, + 'high_prec': True, + 'type_sel': None, + 'repeat': 1, + 'default': 0.0}, + 'force': {'ndof': 3, + 'atomic': True, + 'must': False, + 'high_prec': False, + 'type_sel': None, + 'repeat': 1, + 'default': 0.0}, + 'virial': {'ndof': 9, + 'atomic': False, + 'must': False, + 'high_prec': False, + 'type_sel': None, + 'repeat': 1, + 'default': 0.0}, + 'atom_ener': {'ndof': 1, + 'atomic': True, + 'must': False, + 'high_prec': False, + 'type_sel': None, + 'repeat': 1, + 'default': 0.0}, + 'atom_pref': {'ndof': 1, + 'atomic': True, + 'must': False, + 'high_prec': False, + 'type_sel': None, + 'repeat': 3, + 'default': 0.0}} + data.add_dict(data_requirement) + stop_batch = jdata["training"]["numb_steps"] + + return INPUT, ckpt, frozen_model, model_ckpt, model_frz, data, stop_batch + + +INPUT, CKPT, FROZEN_MODEL, CKPT_TRAINER, FRZ_TRAINER, VALID_DATA, STOP_BATCH = _init_models() + + +class TestInitFrzModelR(unittest.TestCase): + @classmethod + def setUpClass(self): + self.dp_ckpt = CKPT_TRAINER + self.dp_frz = FRZ_TRAINER + self.valid_data = VALID_DATA + self.stop_batch = STOP_BATCH + + @classmethod + def tearDownClass(self): + _file_delete(INPUT) + _file_delete(FROZEN_MODEL) + _file_delete("out.json") + _file_delete(str(tests_path / "checkpoint")) + _file_delete(CKPT+".meta") + _file_delete(CKPT+".index") + _file_delete(CKPT+".data-00000-of-00001") + _file_delete(CKPT+"-0.meta") + _file_delete(CKPT+"-0.index") + _file_delete(CKPT+"-0.data-00000-of-00001") + _file_delete(CKPT+"-1.meta") + _file_delete(CKPT+"-1.index") + _file_delete(CKPT+"-1.data-00000-of-00001") + _file_delete("input_v2_compat.json") + _file_delete("lcurve.out") + + def test_single_frame(self): + valid_batch = self.valid_data.get_batch() + natoms = valid_batch['natoms_vec'] + tf.reset_default_graph() + self.dp_ckpt.build(self.valid_data, self.stop_batch) + self.dp_ckpt._init_session() + feed_dict_ckpt = self.dp_ckpt.get_feed_dict(valid_batch, is_training=False) + ckpt_rmse_ckpt = self.dp_ckpt.loss.eval(self.dp_ckpt.sess, feed_dict_ckpt, natoms) + tf.reset_default_graph() + + self.dp_frz.build(self.valid_data, self.stop_batch) + self.dp_frz._init_session() + feed_dict_frz = self.dp_frz.get_feed_dict(valid_batch, is_training=False) + ckpt_rmse_frz = self.dp_frz.loss.eval(self.dp_frz.sess, feed_dict_frz, natoms) + tf.reset_default_graph() + + # check values + np.testing.assert_almost_equal(ckpt_rmse_ckpt['rmse_e'], ckpt_rmse_frz['rmse_e'], default_places) + np.testing.assert_almost_equal(ckpt_rmse_ckpt['rmse_f'], ckpt_rmse_frz['rmse_f'], default_places) + np.testing.assert_almost_equal(ckpt_rmse_ckpt['rmse_v'], ckpt_rmse_frz['rmse_v'], default_places)