diff --git a/deepmd/descriptor/se_atten.py b/deepmd/descriptor/se_atten.py index d5a2e565c3..61012cea7e 100644 --- a/deepmd/descriptor/se_atten.py +++ b/deepmd/descriptor/se_atten.py @@ -1,3 +1,4 @@ +import logging import warnings from typing import ( List, @@ -12,6 +13,7 @@ from deepmd.common import ( cast_precision, + get_np_precision, ) from deepmd.env import ( GLOBAL_NP_FLOAT_PRECISION, @@ -23,6 +25,9 @@ ) from deepmd.utils.graph import ( get_attention_layer_variables_from_graph_def, + get_pattern_nodes_from_graph_def, + get_tensor_by_name_from_graph, + get_tensor_by_type, ) from deepmd.utils.network import ( embedding_net, @@ -31,6 +36,9 @@ from deepmd.utils.sess import ( run_sess, ) +from deepmd.utils.tabulate import ( + DPTabulate, +) from .descriptor import ( Descriptor, @@ -39,6 +47,8 @@ DescrptSeA, ) +log = logging.getLogger(__name__) + @Descriptor.register("se_atten") class DescrptSeAtten(DescrptSeA): @@ -110,6 +120,7 @@ def __init__( attn_dotr: bool = True, attn_mask: bool = False, multi_task: bool = False, + stripped_type_embedding: bool = False, **kwargs, ) -> None: if not set_davg_zero: @@ -141,11 +152,15 @@ def __init__( assert Version(TF_VERSION) > Version( "2" ), "se_atten only support tensorflow version 2.0 or higher." + self.stripped_type_embedding = stripped_type_embedding self.ntypes = ntypes self.att_n = attn self.attn_layer = attn_layer self.attn_mask = attn_mask self.attn_dotr = attn_dotr + self.filter_np_precision = get_np_precision(precision) + self.two_side_embeeding_net_variables = None + self.layer_size = len(neuron) # descrpt config self.sel_all_a = [sel] @@ -285,6 +300,185 @@ def compute_input_stats( self.stat_dict["sumr2"] += sumr2 self.stat_dict["suma2"] += suma2 + def enable_compression( + self, + min_nbor_dist: float, + graph: tf.Graph, + graph_def: tf.GraphDef, + table_extrapolate: float = 5, + table_stride_1: float = 0.01, + table_stride_2: float = 0.1, + check_frequency: int = -1, + suffix: str = "", + ) -> None: + """Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data. + + Parameters + ---------- + min_nbor_dist + The nearest distance between atoms + graph : tf.Graph + The graph of the model + graph_def : tf.GraphDef + The graph_def of the model + table_extrapolate + The scale of model extrapolation + table_stride_1 + The uniform stride of the first table + table_stride_2 + The uniform stride of the second table + check_frequency + The overflow check frequency + suffix : str, optional + The suffix of the scope + """ + # do some checks before the mocel compression process + assert ( + not self.filter_resnet_dt + ), "Model compression error: descriptor resnet_dt must be false!" + for tt in self.exclude_types: + if (tt[0] not in range(self.ntypes)) or (tt[1] not in range(self.ntypes)): + raise RuntimeError( + "exclude types" + + str(tt) + + " must within the number of atomic types " + + str(self.ntypes) + + "!" + ) + if self.ntypes * self.ntypes - len(self.exclude_types) == 0: + raise RuntimeError( + "empty embedding-net are not supported in model compression!" + ) + + for ii in range(len(self.filter_neuron) - 1): + if self.filter_neuron[ii] * 2 != self.filter_neuron[ii + 1]: + raise NotImplementedError( + "Model Compression error: descriptor neuron [%s] is not supported by model compression! " + "The size of the next layer of the neural network must be twice the size of the previous layer." + % ",".join([str(item) for item in self.filter_neuron]) + ) + + if self.attn_layer != 0: + raise RuntimeError("can not compress model when attention layer is not 0.") + + ret = get_pattern_nodes_from_graph_def( + graph_def, f"filter_type_all{suffix}/.+_two_side_ebd" + ) + if len(ret) == 0: + raise RuntimeError( + "can not find variables of embedding net `*_two_side_ebd` from graph_def, maybe it is not a compressible model." + ) + + self.compress = True + self.table = DPTabulate( + self, + self.filter_neuron, + graph, + graph_def, + True, + self.exclude_types, + self.compress_activation_fn, + suffix=suffix, + ) + self.table_config = [ + table_extrapolate, + table_stride_1, + table_stride_2, + check_frequency, + ] + self.lower, self.upper = self.table.build( + min_nbor_dist, table_extrapolate, table_stride_1, table_stride_2 + ) + + self.final_type_embedding = self._get_two_side_type_embedding(graph) + self.matrix = self._get_two_side_embedding_net_variable( + graph_def, "matrix", suffix + ) + self.bias = self._get_two_side_embedding_net_variable(graph_def, "bias", suffix) + self.two_embd = self._make_data(self.final_type_embedding) + + self.davg = get_tensor_by_name_from_graph( + graph, "descrpt_attr%s/t_avg" % suffix + ) + self.dstd = get_tensor_by_name_from_graph( + graph, "descrpt_attr%s/t_std" % suffix + ) + + def _get_two_side_type_embedding(self, graph): + type_embedding = get_tensor_by_name_from_graph(graph, "t_typeebd") + type_embedding = type_embedding.astype(self.filter_np_precision) + type_embedding_shape = type_embedding.shape + type_embedding_nei = np.tile( + np.reshape(type_embedding, [1, type_embedding_shape[0], -1]), + [type_embedding_shape[0], 1, 1], + ) # (ntypes) * ntypes * Y + type_embedding_center = np.tile( + np.reshape(type_embedding, [type_embedding_shape[0], 1, -1]), + [1, type_embedding_shape[0], 1], + ) # ntypes * (ntypes) * Y + two_side_type_embedding = np.concatenate( + [type_embedding_nei, type_embedding_center], -1 + ) # ntypes * ntypes * (Y+Y) + two_side_type_embedding = np.reshape( + two_side_type_embedding, [-1, two_side_type_embedding.shape[-1]] + ) + return two_side_type_embedding + + def _get_two_side_embedding_net_variable(self, graph_def, varialbe_name, suffix): + ret = {} + for i in range(1, self.layer_size + 1): + target = get_pattern_nodes_from_graph_def( + graph_def, + f"filter_type_all{suffix}/{varialbe_name}_{i}{suffix}_two_side_ebd", + ) + node = target[ + f"filter_type_all{suffix}/{varialbe_name}_{i}{suffix}_two_side_ebd" + ] + ret["layer_" + str(i)] = node + return ret + + def _layer_0(self, x, w, b): + return self.filter_activation_fn(tf.matmul(x, w) + b) + + def _layer_1(self, x, w, b): + t = tf.concat([x, x], axis=1) + return t, self.filter_activation_fn(tf.matmul(x, w) + b) + t + + def _make_data(self, xx): + with tf.Session() as sess: + for layer in range(self.layer_size): + if layer == 0: + if self.filter_neuron[0] == 1: + yy = ( + self._layer_0( + xx, + self.matrix["layer_" + str(layer + 1)], + self.bias["layer_" + str(layer + 1)], + ) + + xx + ) + elif self.filter_neuron[0] == 2: + tt, yy = self._layer_1( + xx, + self.matrix["layer_" + str(layer + 1)], + self.bias["layer_" + str(layer + 1)], + ) + else: + yy = self._layer_0( + xx, + self.matrix["layer_" + str(layer + 1)], + self.bias["layer_" + str(layer + 1)], + ) + else: + tt, zz = self._layer_1( + yy, + self.matrix["layer_" + str(layer + 1)], + self.bias["layer_" + str(layer + 1)], + ) + yy = zz + vv = sess.run(zz) + return vv + def build( self, coord_: tf.Tensor, @@ -806,31 +1000,128 @@ def _filter_lower( # with (natom x nei_type_i) x 1 xyz_scatter = tf.reshape(tf.slice(inputs_reshape, [0, 0], [-1, 1]), [-1, 1]) assert atype is not None, "atype must exist!!" - type_embedding = tf.cast(type_embedding, self.filter_precision) - xyz_scatter = self._lookup_type_embedding(xyz_scatter, atype, type_embedding) - if self.compress: - raise RuntimeError( - "compression of attention descriptor is not supported at the moment" - ) + type_embedding = tf.cast(type_embedding, self.filter_precision) # ntypes * Y # natom x 4 x outputs_size if not is_exclude: with tf.variable_scope(name, reuse=reuse): # with (natom x nei_type_i) x out_size - xyz_scatter = embedding_net( - xyz_scatter, - self.filter_neuron, - self.filter_precision, - activation_fn=activation_fn, - resnet_dt=self.filter_resnet_dt, - name_suffix=suffix, - stddev=stddev, - bavg=bavg, - seed=self.seed, - trainable=trainable, - uniform_seed=self.uniform_seed, - initial_variables=self.embedding_net_variables, - mixed_prec=self.mixed_prec, - ) + if not self.stripped_type_embedding: + log.info("use the previous se_atten model") + xyz_scatter = self._lookup_type_embedding( + xyz_scatter, atype, type_embedding + ) + xyz_scatter = embedding_net( + xyz_scatter, + self.filter_neuron, + self.filter_precision, + activation_fn=activation_fn, + resnet_dt=self.filter_resnet_dt, + name_suffix=suffix, + stddev=stddev, + bavg=bavg, + seed=self.seed, + trainable=trainable, + uniform_seed=self.uniform_seed, + initial_variables=self.embedding_net_variables, + mixed_prec=self.mixed_prec, + ) + else: + if self.attn_layer == 0: + log.info( + "use the compressible model with stripped type embedding" + ) + else: + log.info( + "use the non-compressible model with stripped type embedding" + ) + if not self.compress: + xyz_scatter = embedding_net( + xyz_scatter, + self.filter_neuron, + self.filter_precision, + activation_fn=activation_fn, + resnet_dt=self.filter_resnet_dt, + name_suffix=suffix, + stddev=stddev, + bavg=bavg, + seed=self.seed, + trainable=trainable, + uniform_seed=self.uniform_seed, + initial_variables=self.embedding_net_variables, + mixed_prec=self.mixed_prec, + ) + else: + net = "filter_net" + info = [ + self.lower[net], + self.upper[net], + self.upper[net] * self.table_config[0], + self.table_config[1], + self.table_config[2], + self.table_config[3], + ] + + padding_ntypes = type_embedding.shape[ + 0 + ] # this must be self.ntypes + 1 + atype_expand = tf.reshape(atype, [-1, 1]) + idx_i = tf.tile(atype_expand * padding_ntypes, [1, self.nnei]) + idx_j = tf.reshape(self.nei_type_vec, [-1, self.nnei]) + idx = idx_i + idx_j + index_of_two_side = tf.reshape(idx, [-1]) + + if self.compress: + two_embd = tf.nn.embedding_lookup( + self.two_embd, index_of_two_side + ) + else: + type_embedding_nei = tf.tile( + tf.reshape(type_embedding, [1, padding_ntypes, -1]), + [padding_ntypes, 1, 1], + ) # (ntypes) * ntypes * Y + type_embedding_center = tf.tile( + tf.reshape(type_embedding, [padding_ntypes, 1, -1]), + [1, padding_ntypes, 1], + ) # ntypes * (ntypes) * Y + two_side_type_embedding = tf.concat( + [type_embedding_nei, type_embedding_center], -1 + ) # ntypes * ntypes * (Y+Y) + two_side_type_embedding = tf.reshape( + two_side_type_embedding, + [-1, two_side_type_embedding.shape[-1]], + ) + two_side_type_embedding_suffix = suffix + "_two_side_ebd" + embedding_of_two_side_type_embedding = embedding_net( + two_side_type_embedding, + self.filter_neuron, + self.filter_precision, + activation_fn=activation_fn, + resnet_dt=self.filter_resnet_dt, + name_suffix=two_side_type_embedding_suffix, + stddev=stddev, + bavg=bavg, + seed=self.seed, + trainable=trainable, + uniform_seed=self.uniform_seed, + initial_variables=self.two_side_embeeding_net_variables, + mixed_prec=self.mixed_prec, + ) + two_embd = tf.nn.embedding_lookup( + embedding_of_two_side_type_embedding, index_of_two_side + ) + + if not self.compress: + xyz_scatter = xyz_scatter * two_embd + two_embd + else: + return op_module.tabulate_fusion_se_atten( + tf.cast(self.table.data[net], self.filter_precision), + info, + xyz_scatter, + tf.reshape(inputs_i, [natom, shape_i[1] // 4, 4]), + two_embd, + last_layer_size=outputs_size[-1], + ) + if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift input_r = tf.slice( @@ -961,6 +1252,19 @@ def init_variables( The suffix of the scope """ super().init_variables(graph=graph, graph_def=graph_def, suffix=suffix) + + if self.stripped_type_embedding: + self.two_side_embeeding_net_variables = {} + for i in range(1, self.layer_size + 1): + matrix_pattern = f"filter_type_all{suffix}/matrix_{i}_two_side_ebd" + self.two_side_embeeding_net_variables[ + matrix_pattern + ] = self._get_two_embed_variables(graph_def, matrix_pattern) + bias_pattern = f"filter_type_all{suffix}/bias_{i}_two_side_ebd" + self.two_side_embeeding_net_variables[ + bias_pattern + ] = self._get_two_embed_variables(graph_def, bias_pattern) + self.attention_layer_variables = get_attention_layer_variables_from_graph_def( graph_def, suffix=suffix ) @@ -983,6 +1287,19 @@ def init_variables( ) ] + def _get_two_embed_variables(self, graph_def, pattern: str): + node = get_pattern_nodes_from_graph_def(graph_def, pattern)[pattern] + 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) + return np.reshape(tensor_value, tensor_shape) + def build_type_exclude_mask( self, exclude_types: List[Tuple[int, int]], diff --git a/deepmd/model/ener.py b/deepmd/model/ener.py index 2491091dac..dceac7a0ab 100644 --- a/deepmd/model/ener.py +++ b/deepmd/model/ener.py @@ -368,7 +368,9 @@ def init_variables( else: raise RuntimeError("Unknown model type %s" % model_type) if self.typeebd is not None: - self.typeebd.init_variables(graph, graph_def, suffix=suffix) + self.typeebd.init_variables( + graph, graph_def, suffix=suffix, model_type=model_type + ) def natoms_match(self, force, natoms): use_spin = self.spin.use_spin diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 680d7a9007..74bccbc6a1 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -334,6 +334,7 @@ def descrpt_hybrid_args(): @descrpt_args_plugin.register("se_atten") def descrpt_se_atten_args(): + doc_stripped_type_embedding = "Whether to strip the type embedding into a separated embedding network. Setting it to `False` will fall back to the previous version of `se_atten` which is non-compressible." doc_sel = 'This parameter set the number of selected neighbors. Note that this parameter is a little different from that in other descriptors. Instead of separating each type of atoms, only the summation matters. And this number is highly related with the efficiency, thus one should not make it too large. Usually 200 or less is enough, far away from the GPU limitation 4096. It can be:\n\n\ - `int`. The maximum number of neighbor atoms to be considered. We recommend it to be less than 200. \n\n\ - `List[int]`. The length of the list should be the same as the number of atom types in the system. `sel[i]` gives the selected number of type-i neighbors. Only the summation of `sel[i]` matters, and it is recommended to be less than 200.\ @@ -351,11 +352,18 @@ def descrpt_se_atten_args(): doc_set_davg_zero = "Set the normalization average to zero. This option should be set when `se_atten` descriptor or `atom_ener` in the energy fitting is used" doc_exclude_types = "The excluded pairs of types which have no interaction with each other. For example, `[[0, 1]]` means no interaction between type 0 and type 1." doc_attn = "The length of hidden vectors in attention layers" - doc_attn_layer = "The number of attention layers" + doc_attn_layer = "The number of attention layers. Note that model compression of `se_atten` is only enabled when attn_layer==0 and stripped_type_embedding is True" doc_attn_dotr = "Whether to do dot product with the normalized relative coordinates" doc_attn_mask = "Whether to do mask on the diagonal in the attention matrix" return [ + Argument( + "stripped_type_embedding", + bool, + optional=True, + default=False, + doc=doc_stripped_type_embedding, + ), Argument("sel", [int, list, str], optional=True, default="auto", doc=doc_sel), Argument("rcut", float, optional=True, default=6.0, doc=doc_rcut), Argument("rcut_smth", float, optional=True, default=0.5, doc=doc_rcut_smth), diff --git a/deepmd/utils/tabulate.py b/deepmd/utils/tabulate.py index a31bc7b7f7..6895fcbf13 100644 --- a/deepmd/utils/tabulate.py +++ b/deepmd/utils/tabulate.py @@ -175,7 +175,22 @@ def build( """ # tabulate range [lower, upper] with stride0 'stride0' lower, upper = self._get_env_mat_range(min_nbor_dist) - if isinstance(self.descrpt, deepmd.descriptor.DescrptSeA): + if isinstance(self.descrpt, deepmd.descriptor.DescrptSeAtten): + uu = np.max(upper) + ll = np.min(lower) + xx = np.arange(ll, uu, stride0, dtype=self.data_type) + xx = np.append( + xx, + np.arange(uu, extrapolate * uu, stride1, dtype=self.data_type), + ) + xx = np.append(xx, np.array([extrapolate * uu], dtype=self.data_type)) + nspline = ((uu - ll) / stride0 + (extrapolate * uu - uu) / stride1).astype( + int + ) + self._build_lower( + "filter_net", xx, 0, uu, ll, stride0, stride1, extrapolate, nspline + ) + elif isinstance(self.descrpt, deepmd.descriptor.DescrptSeA): for ii in range(self.table_size): if (self.type_one_side and not self._all_excluded(ii)) or ( not self.type_one_side @@ -403,7 +418,12 @@ def _get_bias(self): bias = {} for layer in range(1, self.layer_size + 1): bias["layer_" + str(layer)] = [] - if isinstance(self.descrpt, deepmd.descriptor.DescrptSeA): + if isinstance(self.descrpt, deepmd.descriptor.DescrptSeAtten): + node = self.embedding_net_nodes[ + f"filter_type_all{self.suffix}/bias_{layer}{self.suffix}" + ] + bias["layer_" + str(layer)].append(tf.make_ndarray(node)) + elif isinstance(self.descrpt, deepmd.descriptor.DescrptSeA): if self.type_one_side: for ii in range(0, self.ntypes): if not self._all_excluded(ii): @@ -462,7 +482,12 @@ def _get_matrix(self): matrix = {} for layer in range(1, self.layer_size + 1): matrix["layer_" + str(layer)] = [] - if isinstance(self.descrpt, deepmd.descriptor.DescrptSeA): + if isinstance(self.descrpt, deepmd.descriptor.DescrptSeAtten): + node = self.embedding_net_nodes[ + f"filter_type_all{self.suffix}/matrix_{layer}{self.suffix}" + ] + matrix["layer_" + str(layer)].append(tf.make_ndarray(node)) + elif isinstance(self.descrpt, deepmd.descriptor.DescrptSeA): if self.type_one_side: for ii in range(0, self.ntypes): if not self._all_excluded(ii): @@ -661,7 +686,9 @@ def _spline5_switch(self, xx, rmin, rmax): def _get_layer_size(self): layer_size = 0 - if isinstance(self.descrpt, deepmd.descriptor.DescrptSeA): + if isinstance(self.descrpt, deepmd.descriptor.DescrptSeAtten): + layer_size = len(self.embedding_net_nodes) // 2 + elif isinstance(self.descrpt, deepmd.descriptor.DescrptSeA): layer_size = len(self.embedding_net_nodes) // ( (self.ntypes * self.ntypes - len(self.exclude_types)) * 2 ) @@ -709,7 +736,9 @@ def _all_excluded(self, ii: int) -> bool: def _get_table_size(self): table_size = 0 - if isinstance(self.descrpt, deepmd.descriptor.DescrptSeA): + if isinstance(self.descrpt, deepmd.descriptor.DescrptSeAtten): + table_size = 1 + elif isinstance(self.descrpt, deepmd.descriptor.DescrptSeA): table_size = self.ntypes * self.ntypes if self.type_one_side: table_size = self.ntypes diff --git a/deepmd/utils/type_embed.py b/deepmd/utils/type_embed.py index 31bd63e405..025747d3b5 100644 --- a/deepmd/utils/type_embed.py +++ b/deepmd/utils/type_embed.py @@ -12,6 +12,7 @@ tf, ) from deepmd.utils.graph import ( + get_tensor_by_name_from_graph, get_type_embedding_net_variables_from_graph_def, ) from deepmd.utils.network import ( @@ -104,7 +105,9 @@ def __init__( self.trainable = trainable self.uniform_seed = uniform_seed self.type_embedding_net_variables = None + self.type_embedding_from_graph = None self.padding = padding + self.model_type = None def build( self, @@ -128,6 +131,8 @@ def build( embedded_types The computational graph for embedded types """ + if self.model_type is not None and self.model_type == "compressed_model": + return self.type_embedding_from_graph types = tf.convert_to_tensor(list(range(ntypes)), dtype=tf.int32) ebd_type = tf.cast( tf.one_hot(tf.cast(types, dtype=tf.int32), int(ntypes)), @@ -159,6 +164,7 @@ def init_variables( graph: tf.Graph, graph_def: tf.GraphDef, suffix="", + model_type="original_model", ) -> None: """Init the type embedding net variables with the given dict. @@ -170,7 +176,12 @@ def init_variables( The input frozen model graph_def suffix Name suffix to identify this descriptor + model_type + Indicator of whether this model is a compressed model """ + self.model_type = model_type self.type_embedding_net_variables = ( get_type_embedding_net_variables_from_graph_def(graph_def, suffix=suffix) ) + type_embedding = get_tensor_by_name_from_graph(graph, "t_typeebd") + self.type_embedding_from_graph = tf.convert_to_tensor(type_embedding) diff --git a/doc/freeze/compress.md b/doc/freeze/compress.md index 4ac45f4de7..6d96c7f01c 100644 --- a/doc/freeze/compress.md +++ b/doc/freeze/compress.md @@ -82,7 +82,7 @@ The model compression interface requires the version of DeePMD-kit used in the o **Acceptable descriptor type** -Descriptors with `se_e2_a`, `se_e3`, and `se_e2_r` types are supported by the model compression feature. `Hybrid` mixed with the above descriptors is also supported. +Descriptors with `se_e2_a`, `se_e3`, `se_e2_r` and `se_atten` types are supported by the model compression feature. `Hybrid` mixed with the above descriptors is also supported. **Available activation functions for descriptor:** diff --git a/examples/water/se_atten/input.json b/examples/water/se_atten/input.json index 3aeeeedbf2..3b8a7bc345 100644 --- a/examples/water/se_atten/input.json +++ b/examples/water/se_atten/input.json @@ -7,6 +7,7 @@ ], "descriptor": { "type": "se_atten", + "stripped_type_embedding": true, "sel": 120, "rcut_smth": 0.50, "rcut": 6.00, @@ -19,7 +20,7 @@ "axis_neuron": 16, "seed": 1, "attn": 128, - "attn_layer": 2, + "attn_layer": 0, "attn_dotr": true, "attn_mask": false, "_comment": " that's all" diff --git a/source/lib/include/tabulate.h b/source/lib/include/tabulate.h index 3ad79f3cb1..7edd3b4d01 100644 --- a/source/lib/include/tabulate.h +++ b/source/lib/include/tabulate.h @@ -8,6 +8,7 @@ void tabulate_fusion_se_a_cpu(FPTYPE* out, const FPTYPE* table_info, const FPTYPE* em_x, const FPTYPE* em, + const FPTYPE* two_embed, const int nloc, const int nnei, const int last_layer_size); @@ -19,6 +20,7 @@ void tabulate_fusion_se_a_grad_cpu(FPTYPE* dy_dem_x, const FPTYPE* table_info, const FPTYPE* em_x, const FPTYPE* em, + const FPTYPE* two_embed, const FPTYPE* dy, const int nloc, const int nnei, @@ -109,6 +111,7 @@ void tabulate_fusion_se_a_gpu_cuda(FPTYPE* out, const FPTYPE* table_info, const FPTYPE* em_x, const FPTYPE* em, + const FPTYPE* two_embed, const int nloc, const int nnei, const int last_layer_size); @@ -120,6 +123,7 @@ void tabulate_fusion_se_a_grad_gpu_cuda(FPTYPE* dy_dem_x, const FPTYPE* table_info, const FPTYPE* em_x, const FPTYPE* em, + const FPTYPE* two_embed, const FPTYPE* dy, const int nloc, const int nnei, @@ -211,6 +215,7 @@ void tabulate_fusion_se_a_gpu_rocm(FPTYPE* out, const FPTYPE* table_info, const FPTYPE* em_x, const FPTYPE* em, + const FPTYPE* two_embed, const int nloc, const int nnei, const int last_layer_size); @@ -222,6 +227,7 @@ void tabulate_fusion_se_a_grad_gpu_rocm(FPTYPE* dy_dem_x, const FPTYPE* table_info, const FPTYPE* em_x, const FPTYPE* em, + const FPTYPE* two_embed, const FPTYPE* dy, const int nloc, const int nnei, diff --git a/source/lib/src/cuda/tabulate.cu b/source/lib/src/cuda/tabulate.cu index 5c3d360c25..a63038f406 100644 --- a/source/lib/src/cuda/tabulate.cu +++ b/source/lib/src/cuda/tabulate.cu @@ -120,6 +120,7 @@ __global__ void tabulate_fusion_se_a_fifth_order_polynomial( const FPTYPE* table, const FPTYPE* em_x, const FPTYPE* em, + const FPTYPE* two_embed, const FPTYPE lower, const FPTYPE upper, const FPTYPE max, @@ -127,6 +128,7 @@ __global__ void tabulate_fusion_se_a_fifth_order_polynomial( const FPTYPE stride1, const int nnei, const int last_layer_size) { + bool enable_se_atten = two_embed != nullptr; const int_64 block_idx = blockIdx.x; // nloc const int thread_idx = threadIdx.x; // last_layer_size FPTYPE ago = __shfl_sync(0xffffffff, em_x[block_idx * nnei + nnei - 1], 0); @@ -152,6 +154,11 @@ __global__ void tabulate_fusion_se_a_fifth_order_polynomial( var[0] + (var[1] + (var[2] + (var[3] + (var[4] + var[5] * xx) * xx) * xx) * xx) * xx; + if (enable_se_atten) { + FPTYPE t = two_embed[block_idx * nnei * last_layer_size + + ii * last_layer_size + thread_idx]; + res = res * t + t; + } for (int kk = 0; kk < MTILE; kk++) { sum[kk] += (nnei - breakpoint) * @@ -173,6 +180,7 @@ __global__ void tabulate_fusion_se_a_grad_fifth_order_polynomial( const FPTYPE* table, const FPTYPE* em_x, const FPTYPE* em, + const FPTYPE* two_embed, const FPTYPE* dy, const FPTYPE lower, const FPTYPE upper, @@ -181,6 +189,7 @@ __global__ void tabulate_fusion_se_a_grad_fifth_order_polynomial( const FPTYPE stride1, const int nnei, const int last_layer_size) { + bool enable_se_atten = two_embed != nullptr; extern __shared__ int _data[]; const int_64 block_idx = blockIdx.x; // nloc const int thread_idx = threadIdx.x; // KTILE * WARP_SIZE, usally 128 here~ @@ -221,6 +230,12 @@ __global__ void tabulate_fusion_se_a_grad_fifth_order_polynomial( (var[1] + (var[2] + (var[3] + (var[4] + var[5] * xx) * xx) * xx) * xx) * xx; + FPTYPE t; + if (enable_se_atten) { + t = two_embed[block_idx * nnei * last_layer_size + + ii * last_layer_size + jj]; + res = res * t + t; + } for (int kk = 0; kk < MTILE; kk++) { sum[kk] += @@ -235,7 +250,7 @@ __global__ void tabulate_fusion_se_a_grad_fifth_order_polynomial( (var[1] + (2 * var[2] + (3 * var[3] + (4 * var[4] + 5 * var[5] * xx) * xx) * xx) * xx) * - res; + (enable_se_atten ? res * t : res); } __syncwarp(); for (int kk = 0; kk < MTILE; kk++) { @@ -600,6 +615,7 @@ void tabulate_fusion_se_a_gpu_cuda(FPTYPE* out, const FPTYPE* table_info, const FPTYPE* em_x, const FPTYPE* em, + const FPTYPE* two_embed, const int nloc, const int nnei, const int last_layer_size) { @@ -607,9 +623,9 @@ void tabulate_fusion_se_a_gpu_cuda(FPTYPE* out, return; } tabulate_fusion_se_a_fifth_order_polynomial - <<>>(out, table, em_x, em, table_info[0], - table_info[1], table_info[2], table_info[3], - table_info[4], nnei, last_layer_size); + <<>>( + out, table, em_x, em, two_embed, table_info[0], table_info[1], + table_info[2], table_info[3], table_info[4], nnei, last_layer_size); DPErrcheck(cudaGetLastError()); DPErrcheck(cudaDeviceSynchronize()); } @@ -621,6 +637,7 @@ void tabulate_fusion_se_a_grad_gpu_cuda(FPTYPE* dy_dem_x, const FPTYPE* table_info, const FPTYPE* em_x, const FPTYPE* em, + const FPTYPE* two_embed, const FPTYPE* dy, const int nloc, const int nnei, @@ -633,8 +650,9 @@ void tabulate_fusion_se_a_grad_gpu_cuda(FPTYPE* dy_dem_x, tabulate_fusion_se_a_grad_fifth_order_polynomial <<>>( - dy_dem_x, dy_dem, table, em_x, em, dy, table_info[0], table_info[1], - table_info[2], table_info[3], table_info[4], nnei, last_layer_size); + dy_dem_x, dy_dem, table, em_x, em, two_embed, dy, table_info[0], + table_info[1], table_info[2], table_info[3], table_info[4], nnei, + last_layer_size); DPErrcheck(cudaGetLastError()); DPErrcheck(cudaDeviceSynchronize()); } @@ -805,6 +823,7 @@ template void tabulate_fusion_se_a_gpu_cuda(float* out, const float* table_info, const float* em_x, const float* em, + const float* two_embed, const int nloc, const int nnei, const int last_layer_size); @@ -813,6 +832,7 @@ template void tabulate_fusion_se_a_gpu_cuda(double* out, const double* table_info, const double* em_x, const double* em, + const double* two_embed, const int nloc, const int nnei, const int last_layer_size); @@ -823,6 +843,7 @@ template void tabulate_fusion_se_a_grad_gpu_cuda( const float* table_info, const float* em_x, const float* em, + const float* two_embed, const float* dy, const int nloc, const int nnei, @@ -834,6 +855,7 @@ template void tabulate_fusion_se_a_grad_gpu_cuda( const double* table_info, const double* em_x, const double* em, + const double* two_embed, const double* dy, const int nloc, const int nnei, diff --git a/source/lib/src/rocm/tabulate.hip.cu b/source/lib/src/rocm/tabulate.hip.cu index b356cb6f3e..3f91854643 100644 --- a/source/lib/src/rocm/tabulate.hip.cu +++ b/source/lib/src/rocm/tabulate.hip.cu @@ -80,6 +80,7 @@ __global__ void tabulate_fusion_se_a_fifth_order_polynomial( const FPTYPE* table, const FPTYPE* em_x, const FPTYPE* em, + const FPTYPE* two_embed, const FPTYPE lower, const FPTYPE upper, const FPTYPE max, @@ -87,6 +88,7 @@ __global__ void tabulate_fusion_se_a_fifth_order_polynomial( const FPTYPE stride1, const int nnei, const int last_layer_size) { + bool enable_se_atten = two_embed != nullptr; HIP_DYNAMIC_SHARED(int, _data) const int_64 block_idx = blockIdx.x; // nloc const int thread_idx = threadIdx.x; // last_layer_size @@ -117,6 +119,11 @@ __global__ void tabulate_fusion_se_a_fifth_order_polynomial( var[0] + (var[1] + (var[2] + (var[3] + (var[4] + var[5] * xx) * xx) * xx) * xx) * xx; + if (enable_se_atten) { + FPTYPE t = two_embed[block_idx * nnei * last_layer_size + + ii * last_layer_size + thread_idx]; + res = res * t + t; + } for (int kk = 0; kk < MTILE; kk++) { iteratorC[kk * last_layer_size + thread_idx] += @@ -138,6 +145,7 @@ __global__ void tabulate_fusion_se_a_grad_fifth_order_polynomial( const FPTYPE* table, const FPTYPE* em_x, const FPTYPE* em, + const FPTYPE* two_embed, const FPTYPE* dy, const FPTYPE lower, const FPTYPE upper, @@ -146,6 +154,7 @@ __global__ void tabulate_fusion_se_a_grad_fifth_order_polynomial( const FPTYPE stride1, const int nnei, const int last_layer_size) { + bool enable_se_atten = two_embed != nullptr; HIP_DYNAMIC_SHARED(int, _data) const int_64 block_idx = blockIdx.x; // nloc const int thread_idx = threadIdx.x; // KTILE * WARP_SIZE, usally 128 here~ @@ -187,6 +196,12 @@ __global__ void tabulate_fusion_se_a_grad_fifth_order_polynomial( (var[1] + (var[2] + (var[3] + (var[4] + var[5] * xx) * xx) * xx) * xx) * xx; + FPTYPE t; + if (enable_se_atten) { + t = two_embed[block_idx * nnei * last_layer_size + + ii * last_layer_size + jj]; + res = res * t + t; + } for (int kk = 0; kk < KTILE; kk++) { sum[kk] += @@ -207,7 +222,7 @@ __global__ void tabulate_fusion_se_a_grad_fifth_order_polynomial( ((FPTYPE)4. * var[4] + (FPTYPE)5. * var[5] * xx) * xx) * xx) * xx) * - res; + (enable_se_atten ? res * t : res); } //__syncwarp();->syncwrap __syncthreads(); @@ -599,6 +614,7 @@ void tabulate_fusion_se_a_gpu_rocm(FPTYPE* out, const FPTYPE* table_info, const FPTYPE* em_x, const FPTYPE* em, + const FPTYPE* two_embed, const int nloc, const int nnei, const int last_layer_size) { @@ -609,7 +625,7 @@ void tabulate_fusion_se_a_gpu_rocm(FPTYPE* out, HIP_KERNEL_NAME( tabulate_fusion_se_a_fifth_order_polynomial), nloc, last_layer_size, sizeof(FPTYPE) * MM * last_layer_size, 0, out, - table, em_x, em, table_info[0], table_info[1], table_info[2], + table, em_x, em, two_embed, table_info[0], table_info[1], table_info[2], table_info[3], table_info[4], nnei, last_layer_size); DPErrcheck(hipGetLastError()); DPErrcheck(hipDeviceSynchronize()); @@ -622,6 +638,7 @@ void tabulate_fusion_se_a_grad_gpu_rocm(FPTYPE* dy_dem_x, const FPTYPE* table_info, const FPTYPE* em_x, const FPTYPE* em, + const FPTYPE* two_embed, const FPTYPE* dy, const int nloc, const int nnei, @@ -636,8 +653,8 @@ void tabulate_fusion_se_a_grad_gpu_rocm(FPTYPE* dy_dem_x, HIP_KERNEL_NAME( tabulate_fusion_se_a_grad_fifth_order_polynomial), nloc, KK * WARP_SIZE, sizeof(FPTYPE) * MM * last_layer_size, 0, dy_dem_x, - dy_dem, table, em_x, em, dy, table_info[0], table_info[1], table_info[2], - table_info[3], table_info[4], nnei, last_layer_size); + dy_dem, table, em_x, em, two_embed, dy, table_info[0], table_info[1], + table_info[2], table_info[3], table_info[4], nnei, last_layer_size); DPErrcheck(hipGetLastError()); DPErrcheck(hipDeviceSynchronize()); } @@ -821,6 +838,7 @@ template void tabulate_fusion_se_a_gpu_rocm(float* out, const float* table_info, const float* em_x, const float* em, + const float* two_embed, const int nloc, const int nnei, const int last_layer_size); @@ -829,6 +847,7 @@ template void tabulate_fusion_se_a_gpu_rocm(double* out, const double* table_info, const double* em_x, const double* em, + const double* two_embed, const int nloc, const int nnei, const int last_layer_size); @@ -839,6 +858,7 @@ template void tabulate_fusion_se_a_grad_gpu_rocm( const float* table_info, const float* em_x, const float* em, + const float* two_embed, const float* dy, const int nloc, const int nnei, @@ -850,6 +870,7 @@ template void tabulate_fusion_se_a_grad_gpu_rocm( const double* table_info, const double* em_x, const double* em, + const double* two_embed, const double* dy, const int nloc, const int nnei, diff --git a/source/lib/src/tabulate.cc b/source/lib/src/tabulate.cc index 67921d4cc7..b2edb30be9 100644 --- a/source/lib/src/tabulate.cc +++ b/source/lib/src/tabulate.cc @@ -82,9 +82,11 @@ void deepmd::tabulate_fusion_se_a_cpu(FPTYPE* out, const FPTYPE* table_info, const FPTYPE* em_x, const FPTYPE* em, + const FPTYPE* two_embed, const int nloc, const int nnei, const int last_layer_size) { + bool enable_se_atten = two_embed != nullptr; memset(out, 0, sizeof(FPTYPE) * nloc * 4 * last_layer_size); const FPTYPE lower = table_info[0]; const FPTYPE upper = table_info[1]; @@ -118,6 +120,12 @@ void deepmd::tabulate_fusion_se_a_cpu(FPTYPE* out, FPTYPE a5 = table[table_idx * last_layer_size * 6 + 6 * kk + 5]; FPTYPE var = a0 + (a1 + (a2 + (a3 + (a4 + a5 * xx) * xx) * xx) * xx) * xx; + if (enable_se_atten) { + FPTYPE t = two_embed[ii * nnei * last_layer_size + + jj * last_layer_size + kk]; + var = var * t + t; + } + if (unloop) { out[ii * last_layer_size * 4 + 0 * last_layer_size + kk] += (nnei - jj) * var * ll[0]; @@ -150,10 +158,12 @@ void deepmd::tabulate_fusion_se_a_grad_cpu(FPTYPE* dy_dem_x, const FPTYPE* table_info, const FPTYPE* em_x, const FPTYPE* em, + const FPTYPE* two_embed, const FPTYPE* dy, const int nloc, const int nnei, const int last_layer_size) { + bool enable_se_atten = two_embed != nullptr; memset(dy_dem_x, 0, sizeof(FPTYPE) * nloc * nnei); memset(dy_dem, 0, sizeof(FPTYPE) * nloc * nnei * 4); FPTYPE const lower = table_info[0]; @@ -195,19 +205,23 @@ void deepmd::tabulate_fusion_se_a_grad_cpu(FPTYPE* dy_dem_x, FPTYPE a5 = table[table_idx * last_layer_size * 6 + 6 * kk + 5]; FPTYPE res = a0 + (a1 + (a2 + (a3 + (a4 + a5 * xx) * xx) * xx) * xx) * xx; + FPTYPE g = + (a1 + (2 * a2 + (3 * a3 + (4 * a4 + 5 * a5 * xx) * xx) * xx) * xx); + if (enable_se_atten) { + FPTYPE t = two_embed[ii * nnei * last_layer_size + + jj * last_layer_size + kk]; + res = res * t + t; + g = t * g; + } if (unloop) { - grad += (a1 + (2 * a2 + (3 * a3 + (4 * a4 + 5 * a5 * xx) * xx) * xx) * - xx) * - dot(ll, rr) * (nnei - jj); + grad += g * dot(ll, rr) * (nnei - jj); dy_dem[ii * nnei * 4 + jj * 4 + 0] += res * rr[0] * (nnei - jj); dy_dem[ii * nnei * 4 + jj * 4 + 1] += res * rr[1] * (nnei - jj); dy_dem[ii * nnei * 4 + jj * 4 + 2] += res * rr[2] * (nnei - jj); dy_dem[ii * nnei * 4 + jj * 4 + 3] += res * rr[3] * (nnei - jj); } else { - grad += (a1 + (2 * a2 + (3 * a3 + (4 * a4 + 5 * a5 * xx) * xx) * xx) * - xx) * - dot(ll, rr); + grad += g * dot(ll, rr); dy_dem[ii * nnei * 4 + jj * 4 + 0] += res * rr[0]; dy_dem[ii * nnei * 4 + jj * 4 + 1] += res * rr[1]; dy_dem[ii * nnei * 4 + jj * 4 + 2] += res * rr[2]; @@ -589,6 +603,7 @@ template void deepmd::tabulate_fusion_se_a_cpu( const float* table_info, const float* em_x, const float* em, + const float* two_embed, const int nloc, const int nnei, const int last_layer_size); @@ -598,6 +613,7 @@ template void deepmd::tabulate_fusion_se_a_cpu( const double* table_info, const double* em_x, const double* em, + const double* two_embed, const int nloc, const int nnei, const int last_layer_size); @@ -608,6 +624,7 @@ template void deepmd::tabulate_fusion_se_a_grad_cpu( const float* table_info, const float* em_x, const float* em, + const float* two_embed, const float* dy, const int nloc, const int nnei, @@ -619,6 +636,7 @@ template void deepmd::tabulate_fusion_se_a_grad_cpu( const double* table_info, const double* em_x, const double* em, + const double* two_embed, const double* dy, const int nloc, const int nnei, diff --git a/source/lib/tests/test_tabulate_se_a.cc b/source/lib/tests/test_tabulate_se_a.cc index f4cf9789cf..984961d4d0 100644 --- a/source/lib/tests/test_tabulate_se_a.cc +++ b/source/lib/tests/test_tabulate_se_a.cc @@ -602,6 +602,94 @@ class TestTabulateSeA : public ::testing::Test { -3.33688909, -3.33688909, -3.34206038, -3.34206038, -3.34206038, -3.34206038, -3.34387412, -3.34387412, -3.34387412, -3.34387412, -3.33051143, -3.33051143, -3.33051143, -3.33051143}; + + std::vector two_embed = { + 0.41783850884461693, 0.06917892522383784, 0.07309949640440838, + 0.57828038123179, 0.30460107001129133, 0.0641857998132136, + 0.016519028000859692, 0.46818914782665344, 0.7524658161955905, + 0.7366050152276675, 0.5442923017739666, 0.6984532784508917, + 0.8794579292532613, 0.933333068809702, 0.052557248156142045, + 0.3158695444821408, 0.6104896498153188, 0.3190616402773879, + 0.39327308944220873, 0.9555810861515368, 0.45845946239660273, + 0.2836952640436372, 0.7129235830370116, 0.21678811087765415, + 0.8589385334305147, 0.8664288996198418, 0.2392088190073245, + 0.44554156483185636, 0.4554902141228184, 0.6929437508125064, + 0.680397459717037, 0.47499455998030615, 0.19872841218252735, + 0.0593083660501722, 0.20593103822290515, 0.8377626566866462, + 0.9006561442856688, 0.9451358048366522, 0.03886827486931199, + 0.9395629463676399, 0.0018941296317954714, 0.08140115779980839, + 0.9309153205767321, 0.4697357303240055, 0.9164471895052549, + 0.5957401092143415, 0.40338864067603986, 0.9096349710860572, + 0.027870073369474335, 0.9614765307496669, 0.3142035164603587, + 0.4027282032956391, 0.05129061735552376, 0.18018240347684844, + 0.8391548601089657, 0.25198333808352436, 0.07903335895654717, + 0.9831396742713334, 0.21328884297544115, 0.8119626188647525, + 0.1734453905861253, 0.014174310472666818, 0.06890574596678134, + 0.3490769935686088, 0.34055562797730554, 0.9834924063503578, + 0.1689164263315952, 0.9611024936313157, 0.6796725725159389, + 0.7902946379060674, 0.3045916985592084, 0.6923776720247495, + 0.5626994287153583, 0.12132066580981216, 0.13356198804830732, + 0.5332034125146011, 0.6155216974624633, 0.3080851791499254, + 0.8391387652641518, 0.8588772315368923, 0.9414859699900482, + 0.9852118289755771, 0.7514252073835589, 0.6780090883007501, + 0.1472721338720271, 0.4785493098407567, 0.8825343095166535, + 0.1562449821247882, 0.5809096109347806, 0.653605647812403, + 0.26158060329219845, 0.28359029181509054, 0.23111396285536823, + 0.6711415141607222, 0.5955230293073148, 0.14336394912405104, + 0.48478135042139503, 0.34621668898158153, 0.7962234329935334, + 0.40204538487553787, 0.09600971949708359, 0.985025266359638, + 0.4949655728846287, 0.23503981206241742, 0.607828476455725, + 0.21634419784756398, 0.04166567958728129, 0.132198384508056, + 0.4112021863641492, 0.9441979803962212, 0.993462871462463, + 0.4524002115880147, 0.6322719605196645, 0.5121196654684579, + 0.7844974619880201, 0.6783684708633317, 0.6402712236722511, + 0.43899788665378925, 0.6060330628471464, 0.7082475921988166, + 0.1614968711069913, 0.6289247345866867, 0.4034261331727077, + 0.7906075239905527, 0.9325509002602962, 0.44489583733770977, + 0.5194672674960213, 0.04635102497306032}; + std::vector expected_xyz_scatter_with_two_embed = { + 0.52052723, -0.09162471, 0.00134297, 0.36123220, 0.53111037, -0.13234888, + 0.14547290, -0.07760225, 0.39022224, -0.05084421, 0.00097584, 0.40299126, + 0.46399123, -0.08226800, 0.10496569, -0.05697074, 0.22831256, -0.03236614, + 0.00053069, 0.20158495, 0.25562350, -0.05008332, 0.05120581, -0.03847130, + 0.48372134, -0.08635551, 0.00127416, 0.33553740, 0.48870492, -0.12377265, + 0.14484348, -0.06894720, 0.09037214, -0.08453476, 0.00132198, 0.48978396, + 0.56809636, -0.11217096, 0.06433977, -0.14684460, 0.04360697, -0.10371710, + 0.00105878, 0.36754591, 0.30793639, -0.05068795, 0.10760094, -0.09443281, + 0.03027943, -0.05996609, 0.00096647, 0.29035073, 0.38261345, -0.05492907, + 0.06442101, -0.09657691, 0.06211007, -0.06585386, 0.00121839, 0.38826136, + 0.62888512, -0.12626045, 0.13760707, -0.15409971, 0.35496658, -0.09629124, + 0.00094936, 0.37091455, 0.40029681, -0.11266203, 0.16359092, -0.09316170, + 0.46220540, -0.11194036, 0.00099500, 0.48276205, 0.60272823, -0.14764563, + 0.23083357, -0.14624559, 0.42737683, -0.12200031, 0.00101439, 0.38477701, + 0.40109773, -0.14840970, 0.18621217, -0.08161032, 0.19972406, -0.03173626, + 0.00020518, 0.11990507, 0.18933921, -0.05445546, 0.10560781, -0.04502732, + 0.43067025, -0.06062120, 0.00243902, 0.29049015, 0.50532163, -0.17136189, + 0.24445441, -0.10595959, 0.24687635, -0.06109584, 0.00155267, 0.31006293, + 0.46913106, -0.14140409, 0.16395873, -0.04099833, 0.18986012, -0.05425738, + 0.00104613, 0.27541918, 0.46324668, -0.11304375, 0.14316380, -0.03016295, + 0.30689686, -0.03886867, 0.00133389, 0.17385121, 0.33724930, -0.13383926, + 0.18766727, -0.07051321, + }; + std::vector expected_dy_dem_x_with_two_embed = { + -0.00876744, -0.05693831, 0.00251697, -0.04104716, + -0.07540836, -0.01255790, -0.00225420, -0.04262001, + -0.09707657, -0.00396457, -0.00492498, -0.03225818, + -0.04350615, -0.02865961, -0.00690901, -0.00847689, + }; + std::vector expected_dy_dem_with_two_embed = { + 1.88405505, 1.88405505, 1.88405505, 1.88405505, 3.33316973, 3.33316973, + 3.33316973, 3.33316973, 3.48744961, 3.48744961, 3.48744961, 3.48744961, + 2.66470385, 2.66470385, 2.66470385, 2.66470385, 2.55291404, 2.55291404, + 2.55291404, 2.55291404, 2.23471214, 2.23471214, 2.23471214, 2.23471214, + 0.79636719, 0.79636719, 0.79636719, 0.79636719, 1.16809199, 1.16809199, + 1.16809199, 1.16809199, 2.55088185, 2.55088185, 2.55088185, 2.55088185, + 3.02601696, 3.02601696, 3.02601696, 3.02601696, 2.80271925, 2.80271925, + 2.80271925, 2.80271925, 1.67897971, 1.67897971, 1.67897971, 1.67897971, + 1.20938629, 1.20938629, 1.20938629, 1.20938629, 2.05980066, 2.05980066, + 2.05980066, 2.05980066, 3.10625208, 3.10625208, 3.10625208, 3.10625208, + 3.35169350, 3.35169350, 3.35169350, 3.35169350, + }; const int nloc = 4; const int nnei = 4; const int last_layer_size = 8; @@ -613,13 +701,24 @@ class TestTabulateSeA : public ::testing::Test { TEST_F(TestTabulateSeA, tabulate_fusion_se_a_cpu) { std::vector xyz_scatter(nloc * nnei * last_layer_size); deepmd::tabulate_fusion_se_a_cpu(&xyz_scatter[0], &table[0], &info[0], - &em_x[0], &em[0], nloc, nnei, - last_layer_size); + &em_x[0], &em[0], nullptr, nloc, + nnei, last_layer_size); EXPECT_EQ(xyz_scatter.size(), nloc * nnei * last_layer_size); EXPECT_EQ(xyz_scatter.size(), expected_xyz_scatter.size()); for (int jj = 0; jj < xyz_scatter.size(); ++jj) { EXPECT_LT(fabs(xyz_scatter[jj] - expected_xyz_scatter[jj]), 1e-5); } + + deepmd::tabulate_fusion_se_a_cpu(&xyz_scatter[0], &table[0], &info[0], + &em_x[0], &em[0], &two_embed[0], + nloc, nnei, last_layer_size); + EXPECT_EQ(xyz_scatter.size(), nloc * nnei * last_layer_size); + EXPECT_EQ(xyz_scatter.size(), expected_xyz_scatter_with_two_embed.size()); + + for (int jj = 0; jj < xyz_scatter.size(); ++jj) { + EXPECT_LT(fabs(xyz_scatter[jj] - expected_xyz_scatter_with_two_embed[jj]), + 1e-5); + } } TEST_F(TestTabulateSeA, tabulate_fusion_se_a_grad_cpu) { @@ -627,8 +726,8 @@ TEST_F(TestTabulateSeA, tabulate_fusion_se_a_grad_cpu) { std::vector dy_dem(em.size()); std::vector dy(nloc * nnei * last_layer_size, 1.0); deepmd::tabulate_fusion_se_a_grad_cpu( - &dy_dem_x[0], &dy_dem[0], &table[0], &info[0], &em_x[0], &em[0], &dy[0], - nloc, nnei, last_layer_size); + &dy_dem_x[0], &dy_dem[0], &table[0], &info[0], &em_x[0], &em[0], nullptr, + &dy[0], nloc, nnei, last_layer_size); EXPECT_EQ(dy_dem_x.size(), nloc * nnei); EXPECT_EQ(dy_dem.size(), nloc * nnei * 4); EXPECT_EQ(dy_dem_x.size(), expected_dy_dem_x.size()); @@ -639,6 +738,20 @@ TEST_F(TestTabulateSeA, tabulate_fusion_se_a_grad_cpu) { for (int jj = 0; jj < dy_dem.size(); ++jj) { EXPECT_LT(fabs(dy_dem[jj] - expected_dy_dem[jj]), 1e-5); } + + deepmd::tabulate_fusion_se_a_grad_cpu( + &dy_dem_x[0], &dy_dem[0], &table[0], &info[0], &em_x[0], &em[0], + &two_embed[0], &dy[0], nloc, nnei, last_layer_size); + EXPECT_EQ(dy_dem_x.size(), nloc * nnei); + EXPECT_EQ(dy_dem.size(), nloc * nnei * 4); + EXPECT_EQ(dy_dem_x.size(), expected_dy_dem_x.size()); + EXPECT_EQ(dy_dem.size(), expected_dy_dem.size()); + for (int jj = 0; jj < dy_dem_x.size(); ++jj) { + EXPECT_LT(fabs(dy_dem_x[jj] - expected_dy_dem_x_with_two_embed[jj]), 1e-5); + } + for (int jj = 0; jj < dy_dem.size(); ++jj) { + EXPECT_LT(fabs(dy_dem[jj] - expected_dy_dem_with_two_embed[jj]), 1e-5); + } } #if GOOGLE_CUDA @@ -651,20 +764,37 @@ TEST_F(TestTabulateSeA, tabulate_fusion_se_a_gpu_cuda) { deepmd::malloc_device_memory_sync(table_dev, table); deepmd::malloc_device_memory_sync(em_x_dev, em_x); deepmd::malloc_device_memory_sync(em_dev, em); - deepmd::tabulate_fusion_se_a_gpu_cuda(xyz_scatter_dev, table_dev, - &info[0], em_x_dev, em_dev, - nloc, nnei, last_layer_size); + deepmd::tabulate_fusion_se_a_gpu_cuda( + xyz_scatter_dev, table_dev, &info[0], em_x_dev, em_dev, nullptr, nloc, + nnei, last_layer_size); deepmd::memcpy_device_to_host(xyz_scatter_dev, xyz_scatter); - deepmd::delete_device_memory(xyz_scatter_dev); - deepmd::delete_device_memory(table_dev); - deepmd::delete_device_memory(em_x_dev); - deepmd::delete_device_memory(em_dev); EXPECT_EQ(xyz_scatter.size(), nloc * nnei * last_layer_size); EXPECT_EQ(xyz_scatter.size(), expected_xyz_scatter.size()); for (int jj = 0; jj < xyz_scatter.size(); ++jj) { EXPECT_LT(fabs(xyz_scatter[jj] - expected_xyz_scatter[jj]), 1e-5); } + + double *two_embed_dev = nullptr; + deepmd::malloc_device_memory_sync(two_embed_dev, two_embed); + deepmd::malloc_device_memory_sync(xyz_scatter_dev, xyz_scatter); + deepmd::tabulate_fusion_se_a_gpu_cuda( + xyz_scatter_dev, table_dev, &info[0], em_x_dev, em_dev, two_embed_dev, + nloc, nnei, last_layer_size); + deepmd::memcpy_device_to_host(xyz_scatter_dev, xyz_scatter); + + EXPECT_EQ(xyz_scatter.size(), nloc * nnei * last_layer_size); + EXPECT_EQ(xyz_scatter.size(), expected_xyz_scatter.size()); + for (int jj = 0; jj < xyz_scatter.size(); ++jj) { + EXPECT_LT(fabs(xyz_scatter[jj] - expected_xyz_scatter_with_two_embed[jj]), + 1e-5); + } + + deepmd::delete_device_memory(xyz_scatter_dev); + deepmd::delete_device_memory(table_dev); + deepmd::delete_device_memory(em_x_dev); + deepmd::delete_device_memory(em_dev); + deepmd::delete_device_memory(two_embed_dev); } TEST_F(TestTabulateSeA, tabulate_fusion_se_a_grad_gpu_cuda) { @@ -681,16 +811,10 @@ TEST_F(TestTabulateSeA, tabulate_fusion_se_a_grad_gpu_cuda) { deepmd::malloc_device_memory_sync(em_dev, em); deepmd::malloc_device_memory_sync(dy_dev, dy); deepmd::tabulate_fusion_se_a_grad_gpu_cuda( - dy_dem_x_dev, dy_dem_dev, table_dev, &info[0], em_x_dev, em_dev, dy_dev, - nloc, nnei, last_layer_size); + dy_dem_x_dev, dy_dem_dev, table_dev, &info[0], em_x_dev, em_dev, nullptr, + dy_dev, nloc, nnei, last_layer_size); deepmd::memcpy_device_to_host(dy_dem_x_dev, dy_dem_x); deepmd::memcpy_device_to_host(dy_dem_dev, dy_dem); - deepmd::delete_device_memory(dy_dem_x_dev); - deepmd::delete_device_memory(dy_dem_dev); - deepmd::delete_device_memory(table_dev); - deepmd::delete_device_memory(em_x_dev); - deepmd::delete_device_memory(em_dev); - deepmd::delete_device_memory(dy_dev); EXPECT_EQ(dy_dem_x.size(), nloc * nnei); EXPECT_EQ(dy_dem.size(), nloc * nnei * 4); @@ -702,6 +826,30 @@ TEST_F(TestTabulateSeA, tabulate_fusion_se_a_grad_gpu_cuda) { for (int jj = 0; jj < dy_dem.size(); ++jj) { EXPECT_LT(fabs(dy_dem[jj] - expected_dy_dem[jj]), 1e-5); } + + double *two_embed_dev = nullptr; + deepmd::malloc_device_memory_sync(two_embed_dev, two_embed); + deepmd::malloc_device_memory_sync(dy_dem_x_dev, dy_dem_x); + deepmd::malloc_device_memory_sync(dy_dem_dev, dy_dem); + deepmd::tabulate_fusion_se_a_grad_gpu_cuda( + dy_dem_x_dev, dy_dem_dev, table_dev, &info[0], em_x_dev, em_dev, + two_embed_dev, dy_dev, nloc, nnei, last_layer_size); + deepmd::memcpy_device_to_host(dy_dem_x_dev, dy_dem_x); + deepmd::memcpy_device_to_host(dy_dem_dev, dy_dem); + for (int jj = 0; jj < dy_dem_x.size(); ++jj) { + EXPECT_LT(fabs(dy_dem_x[jj] - expected_dy_dem_x_with_two_embed[jj]), 1e-5); + } + for (int jj = 0; jj < dy_dem.size(); ++jj) { + EXPECT_LT(fabs(dy_dem[jj] - expected_dy_dem_with_two_embed[jj]), 1e-5); + } + + deepmd::delete_device_memory(dy_dem_x_dev); + deepmd::delete_device_memory(dy_dem_dev); + deepmd::delete_device_memory(table_dev); + deepmd::delete_device_memory(em_x_dev); + deepmd::delete_device_memory(em_dev); + deepmd::delete_device_memory(dy_dev); + deepmd::delete_device_memory(two_embed_dev); } #endif // GOOGLE_CUDA @@ -715,20 +863,37 @@ TEST_F(TestTabulateSeA, tabulate_fusion_se_a_gpu_rocm) { deepmd::malloc_device_memory_sync(table_dev, table); deepmd::malloc_device_memory_sync(em_x_dev, em_x); deepmd::malloc_device_memory_sync(em_dev, em); - deepmd::tabulate_fusion_se_a_gpu_rocm(xyz_scatter_dev, table_dev, - &info[0], em_x_dev, em_dev, - nloc, nnei, last_layer_size); + deepmd::tabulate_fusion_se_a_gpu_rocm( + xyz_scatter_dev, table_dev, &info[0], em_x_dev, em_dev, nullptr, nloc, + nnei, last_layer_size); deepmd::memcpy_device_to_host(xyz_scatter_dev, xyz_scatter); - deepmd::delete_device_memory(xyz_scatter_dev); - deepmd::delete_device_memory(table_dev); - deepmd::delete_device_memory(em_x_dev); - deepmd::delete_device_memory(em_dev); EXPECT_EQ(xyz_scatter.size(), nloc * nnei * last_layer_size); EXPECT_EQ(xyz_scatter.size(), expected_xyz_scatter.size()); for (int jj = 0; jj < xyz_scatter.size(); ++jj) { EXPECT_LT(fabs(xyz_scatter[jj] - expected_xyz_scatter[jj]), 1e-5); } + + double *two_embed_dev = nullptr; + deepmd::malloc_device_memory_sync(two_embed_dev, two_embed); + deepmd::malloc_device_memory_sync(xyz_scatter_dev, xyz_scatter); + deepmd::tabulate_fusion_se_a_gpu_rocm( + xyz_scatter_dev, table_dev, &info[0], em_x_dev, em_dev, two_embed_dev, + nloc, nnei, last_layer_size); + deepmd::memcpy_device_to_host(xyz_scatter_dev, xyz_scatter); + + EXPECT_EQ(xyz_scatter.size(), nloc * nnei * last_layer_size); + EXPECT_EQ(xyz_scatter.size(), expected_xyz_scatter.size()); + for (int jj = 0; jj < xyz_scatter.size(); ++jj) { + EXPECT_LT(fabs(xyz_scatter[jj] - expected_xyz_scatter_with_two_embed[jj]), + 1e-5); + } + + deepmd::delete_device_memory(xyz_scatter_dev); + deepmd::delete_device_memory(table_dev); + deepmd::delete_device_memory(em_x_dev); + deepmd::delete_device_memory(em_dev); + deepmd::delete_device_memory(two_embed_dev); } TEST_F(TestTabulateSeA, tabulate_fusion_se_a_grad_gpu_rocm) { @@ -745,16 +910,10 @@ TEST_F(TestTabulateSeA, tabulate_fusion_se_a_grad_gpu_rocm) { deepmd::malloc_device_memory_sync(em_dev, em); deepmd::malloc_device_memory_sync(dy_dev, dy); deepmd::tabulate_fusion_se_a_grad_gpu_rocm( - dy_dem_x_dev, dy_dem_dev, table_dev, &info[0], em_x_dev, em_dev, dy_dev, - nloc, nnei, last_layer_size); + dy_dem_x_dev, dy_dem_dev, table_dev, &info[0], em_x_dev, em_dev, nullptr, + dy_dev, nloc, nnei, last_layer_size); deepmd::memcpy_device_to_host(dy_dem_x_dev, dy_dem_x); deepmd::memcpy_device_to_host(dy_dem_dev, dy_dem); - deepmd::delete_device_memory(dy_dem_x_dev); - deepmd::delete_device_memory(dy_dem_dev); - deepmd::delete_device_memory(table_dev); - deepmd::delete_device_memory(em_x_dev); - deepmd::delete_device_memory(em_dev); - deepmd::delete_device_memory(dy_dev); EXPECT_EQ(dy_dem_x.size(), nloc * nnei); EXPECT_EQ(dy_dem.size(), nloc * nnei * 4); @@ -766,5 +925,29 @@ TEST_F(TestTabulateSeA, tabulate_fusion_se_a_grad_gpu_rocm) { for (int jj = 0; jj < dy_dem.size(); ++jj) { EXPECT_LT(fabs(dy_dem[jj] - expected_dy_dem[jj]), 1e-5); } + + double *two_embed_dev = nullptr; + deepmd::malloc_device_memory_sync(two_embed_dev, two_embed); + deepmd::malloc_device_memory_sync(dy_dem_x_dev, dy_dem_x); + deepmd::malloc_device_memory_sync(dy_dem_dev, dy_dem); + deepmd::tabulate_fusion_se_a_grad_gpu_rocm( + dy_dem_x_dev, dy_dem_dev, table_dev, &info[0], em_x_dev, em_dev, + two_embed_dev, dy_dev, nloc, nnei, last_layer_size); + deepmd::memcpy_device_to_host(dy_dem_x_dev, dy_dem_x); + deepmd::memcpy_device_to_host(dy_dem_dev, dy_dem); + for (int jj = 0; jj < dy_dem_x.size(); ++jj) { + EXPECT_LT(fabs(dy_dem_x[jj] - expected_dy_dem_x_with_two_embed[jj]), 1e-5); + } + for (int jj = 0; jj < dy_dem.size(); ++jj) { + EXPECT_LT(fabs(dy_dem[jj] - expected_dy_dem_with_two_embed[jj]), 1e-5); + } + + deepmd::delete_device_memory(dy_dem_x_dev); + deepmd::delete_device_memory(dy_dem_dev); + deepmd::delete_device_memory(table_dev); + deepmd::delete_device_memory(em_x_dev); + deepmd::delete_device_memory(em_dev); + deepmd::delete_device_memory(dy_dev); + deepmd::delete_device_memory(two_embed_dev); } #endif // TENSORFLOW_USE_ROCM diff --git a/source/op/_tabulate_grad.py b/source/op/_tabulate_grad.py index 785146cfcc..0d6ef88e3a 100644 --- a/source/op/_tabulate_grad.py +++ b/source/op/_tabulate_grad.py @@ -16,7 +16,12 @@ @ops.RegisterGradient("TabulateFusionSeA") def _tabulate_fusion_se_a_grad_cc(op, dy): dy_dx, dy_df = op_module.tabulate_fusion_se_a_grad( - op.inputs[0], op.inputs[1], op.inputs[2], op.inputs[3], dy, op.outputs[0] + op.inputs[0], + op.inputs[1], + op.inputs[2], + op.inputs[3], + dy, + op.outputs[0], ) return [None, None, dy_dx, dy_df] @@ -30,6 +35,28 @@ def _tabulate_fusion_se_a_grad_grad_cc(op, dy, dy_): return [None, None, None, None, dz_dy, None] +@ops.RegisterGradient("TabulateFusionSeAtten") +def _tabulate_fusion_se_atten_grad_cc(op, dy): + dy_dx, dy_df, dy_dtwo = op_module.tabulate_fusion_se_atten_grad( + op.inputs[0], + op.inputs[1], + op.inputs[2], + op.inputs[3], + op.inputs[4], + dy, + op.outputs[0], + ) + return [None, None, dy_dx, dy_df, dy_dtwo] + + +@ops.RegisterGradient("TabulateFusionSeAttenGrad") +def _tabulate_fusion_se_atten_grad_grad_cc(op, dy, dy_, dy_dtwo): + dz_dy = op_module.tabulate_fusion_se_a_grad_grad( + op.inputs[0], op.inputs[1], op.inputs[2], op.inputs[3], dy, dy_, op.inputs[6] + ) + return [None, None, None, None, None, dz_dy, None] + + @ops.RegisterGradient("TabulateFusionSeT") def _tabulate_fusion_se_t_grad_cc(op, dy): dy_dx, dy_df = op_module.tabulate_fusion_se_t_grad( diff --git a/source/op/tabulate_multi_device.cc b/source/op/tabulate_multi_device.cc index 2914b75c64..09a31f137e 100644 --- a/source/op/tabulate_multi_device.cc +++ b/source/op/tabulate_multi_device.cc @@ -1,3 +1,4 @@ + #include "custom_op.h" #include "tabulate.h" @@ -63,6 +64,29 @@ REGISTER_OP("TabulateFusionSeAGradGrad") .Input("descriptor: T") .Output("dz_dy: T"); +REGISTER_OP("TabulateFusionSeAtten") + .Attr("T: {float, double} = DT_DOUBLE") + .Input("table: T") + .Input("table_info: T") + .Input("em_x: T") + .Input("em: T") + .Input("two_embed: T") + .Attr("last_layer_size: int") + .Output("descriptor: T"); + +REGISTER_OP("TabulateFusionSeAttenGrad") + .Attr("T: {float, double} = DT_DOUBLE") + .Input("table: T") + .Input("table_info: T") + .Input("em_x: T") + .Input("em: T") + .Input("two_embed: T") + .Input("dy: T") + .Input("descriptor: T") + .Output("dy_dem_x: T") + .Output("dy_dem: T") + .Output("dy_dtwo: T"); + REGISTER_OP("TabulateFusionSeT") .Attr("T: {float, double} = DT_DOUBLE") .Input("table: T") @@ -163,22 +187,25 @@ class TabulateFusionSeAOp : public OpKernel { const FPTYPE* table_info = table_info_tensor.flat().data(); const FPTYPE* em_x = em_x_tensor.flat().data(); const FPTYPE* em = em_tensor.flat().data(); + const FPTYPE* two_embed = nullptr; const int nloc = em_tensor.shape().dim_size(0); const int nnei = em_tensor.shape().dim_size(1); if (device == "GPU") { #if GOOGLE_CUDA deepmd::tabulate_fusion_se_a_gpu_cuda(descriptor, table, table_info, em_x, - em, nloc, nnei, last_layer_size); + em, two_embed, nloc, nnei, + last_layer_size); #endif // GOOGLE_CUDA #if TENSORFLOW_USE_ROCM deepmd::tabulate_fusion_se_a_gpu_rocm(descriptor, table, table_info, em_x, - em, nloc, nnei, last_layer_size); + em, two_embed, nloc, nnei, + last_layer_size); #endif // TENSORFLOW_USE_ROCM } else if (device == "CPU") { deepmd::tabulate_fusion_se_a_cpu(descriptor, table, table_info, em_x, em, - nloc, nnei, last_layer_size); + two_embed, nloc, nnei, last_layer_size); } } @@ -204,6 +231,7 @@ class TabulateFusionSeAGradOp : public OpKernel { const Tensor& table_info_tensor = context->input(context_input_index++); const Tensor& em_x_tensor = context->input(context_input_index++); const Tensor& em_tensor = context->input(context_input_index++); + const Tensor& dy_tensor = context->input(context_input_index++); const Tensor& descriptor_tensor = context->input(context_input_index++); // set size of the sample @@ -223,11 +251,13 @@ class TabulateFusionSeAGradOp : public OpKernel { // flat the tensors FPTYPE* dy_dem_x = dy_dem_x_tensor->flat().data(); FPTYPE* dy_dem = dy_dem_tensor->flat().data(); + const FPTYPE* descriptor = descriptor_tensor.flat().data(); const FPTYPE* table = table_tensor.flat().data(); const FPTYPE* table_info = table_info_tensor.flat().data(); const FPTYPE* em_x = em_x_tensor.flat().data(); const FPTYPE* em = em_tensor.flat().data(); + const FPTYPE* two_embed = nullptr; const FPTYPE* dy = dy_tensor.flat().data(); const int nloc = em_tensor.shape().dim_size(0); const int nnei = em_tensor.shape().dim_size(1); @@ -235,19 +265,19 @@ class TabulateFusionSeAGradOp : public OpKernel { if (device == "GPU") { #if GOOGLE_CUDA - deepmd::tabulate_fusion_se_a_grad_gpu_cuda(dy_dem_x, dy_dem, table, - table_info, em_x, em, dy, nloc, - nnei, last_layer_size); + deepmd::tabulate_fusion_se_a_grad_gpu_cuda( + dy_dem_x, dy_dem, table, table_info, em_x, em, two_embed, dy, nloc, + nnei, last_layer_size); #endif // GOOGLE_CUDA #if TENSORFLOW_USE_ROCM - deepmd::tabulate_fusion_se_a_grad_gpu_rocm(dy_dem_x, dy_dem, table, - table_info, em_x, em, dy, nloc, - nnei, last_layer_size); + deepmd::tabulate_fusion_se_a_grad_gpu_rocm( + dy_dem_x, dy_dem, table, table_info, em_x, em, two_embed, dy, nloc, + nnei, last_layer_size); #endif // TENSORFLOW_USE_ROCM } else if (device == "CPU") { deepmd::tabulate_fusion_se_a_grad_cpu(dy_dem_x, dy_dem, table, table_info, - em_x, em, dy, nloc, nnei, + em_x, em, two_embed, dy, nloc, nnei, last_layer_size); } } @@ -321,6 +351,156 @@ class TabulateFusionSeAGradGradOp : public OpKernel { std::string device; }; +template +class TabulateFusionSeAttenOp : public OpKernel { + public: + explicit TabulateFusionSeAttenOp(OpKernelConstruction* context) + : OpKernel(context) { + OP_REQUIRES_OK(context, + context->GetAttr("last_layer_size", &last_layer_size)); + } + void Compute(OpKernelContext* context) override { + deepmd::safe_compute( + context, [this](OpKernelContext* context) { this->_Compute(context); }); + } + + void _Compute(OpKernelContext* context) { + // Grab the input tensor + int context_input_index = 0; + const Tensor& table_tensor = context->input(context_input_index++); + const Tensor& table_info_tensor = context->input(context_input_index++); + const Tensor& em_x_tensor = context->input(context_input_index++); + const Tensor& em_tensor = context->input(context_input_index++); + const Tensor& two_embed_tensor = context->input(context_input_index++); + // set size of the sample + OP_REQUIRES(context, (table_tensor.shape().dims() == 2), + errors::InvalidArgument("Dim of table should be 2")); + OP_REQUIRES(context, (em_x_tensor.shape().dims() == 2), + errors::InvalidArgument("Dim of input should be 2")); + OP_REQUIRES(context, (em_tensor.shape().dims() == 3), + errors::InvalidArgument("Dim of input should be 3")); + OP_REQUIRES(context, (two_embed_tensor.shape().dims() == 2), + errors::InvalidArgument("Dim of input should be 2")); + TensorShape descriptor_shape; + descriptor_shape.AddDim(em_tensor.shape().dim_size(0)); + descriptor_shape.AddDim(4); // TODO: be careful here; + descriptor_shape.AddDim(last_layer_size); + int context_output_index = 0; + Tensor* descriptor_tensor = NULL; + OP_REQUIRES_OK(context, context->allocate_output(context_output_index++, + descriptor_shape, + &descriptor_tensor)); + DeviceFunctor()(device, context->eigen_device()); + // flat the tensors + FPTYPE* descriptor = descriptor_tensor->flat().data(); + const FPTYPE* table = table_tensor.flat().data(); + const FPTYPE* table_info = table_info_tensor.flat().data(); + const FPTYPE* em_x = em_x_tensor.flat().data(); + const FPTYPE* em = em_tensor.flat().data(); + const FPTYPE* two_embed = two_embed_tensor.flat().data(); + const int nloc = em_tensor.shape().dim_size(0); + const int nnei = em_tensor.shape().dim_size(1); + + if (device == "GPU") { +#if GOOGLE_CUDA + deepmd::tabulate_fusion_se_a_gpu_cuda(descriptor, table, table_info, em_x, + em, two_embed, nloc, nnei, + last_layer_size); +#endif // GOOGLE_CUDA + +#if TENSORFLOW_USE_ROCM + deepmd::tabulate_fusion_se_a_gpu_rocm(descriptor, table, table_info, em_x, + em, two_embed, nloc, nnei, + last_layer_size); +#endif // TENSORFLOW_USE_ROCM + } else if (device == "CPU") { + deepmd::tabulate_fusion_se_a_cpu(descriptor, table, table_info, em_x, em, + two_embed, nloc, nnei, last_layer_size); + } + } + + private: + int last_layer_size; + std::string device; +}; + +template +class TabulateFusionSeAttenGradOp : public OpKernel { + public: + explicit TabulateFusionSeAttenGradOp(OpKernelConstruction* context) + : OpKernel(context) {} + void Compute(OpKernelContext* context) override { + deepmd::safe_compute( + context, [this](OpKernelContext* context) { this->_Compute(context); }); + } + + void _Compute(OpKernelContext* context) { + // Grab the input tensor + int context_input_index = 0; + const Tensor& table_tensor = context->input(context_input_index++); + const Tensor& table_info_tensor = context->input(context_input_index++); + const Tensor& em_x_tensor = context->input(context_input_index++); + const Tensor& em_tensor = context->input(context_input_index++); + const Tensor& two_embed_tensor = context->input(context_input_index++); + + const Tensor& dy_tensor = context->input(context_input_index++); + const Tensor& descriptor_tensor = context->input(context_input_index++); + // set size of the sample + OP_REQUIRES(context, (dy_tensor.shape().dims() == 3), + errors::InvalidArgument("Dim of table should be 3")); + int context_output_index = 0; + Tensor* dy_dem_x_tensor = NULL; + OP_REQUIRES_OK(context, context->allocate_output(context_output_index++, + em_x_tensor.shape(), + &dy_dem_x_tensor)); + Tensor* dy_dem_tensor = NULL; + OP_REQUIRES_OK(context, + context->allocate_output(context_output_index++, + em_tensor.shape(), &dy_dem_tensor)); + Tensor* dy_dtwo_tensor = NULL; + OP_REQUIRES_OK(context, context->allocate_output(context_output_index++, + two_embed_tensor.shape(), + &dy_dtwo_tensor)); + DeviceFunctor()(device, context->eigen_device()); + + // flat the tensors + FPTYPE* dy_dem_x = dy_dem_x_tensor->flat().data(); + FPTYPE* dy_dem = dy_dem_tensor->flat().data(); + + const FPTYPE* descriptor = descriptor_tensor.flat().data(); + const FPTYPE* table = table_tensor.flat().data(); + const FPTYPE* table_info = table_info_tensor.flat().data(); + const FPTYPE* em_x = em_x_tensor.flat().data(); + const FPTYPE* em = em_tensor.flat().data(); + const FPTYPE* two_embed = two_embed_tensor.flat().data(); + const FPTYPE* dy = dy_tensor.flat().data(); + const int nloc = em_tensor.shape().dim_size(0); + const int nnei = em_tensor.shape().dim_size(1); + const int last_layer_size = descriptor_tensor.shape().dim_size(2); + + if (device == "GPU") { +#if GOOGLE_CUDA + deepmd::tabulate_fusion_se_a_grad_gpu_cuda( + dy_dem_x, dy_dem, table, table_info, em_x, em, two_embed, dy, nloc, + nnei, last_layer_size); +#endif // GOOGLE_CUDA + +#if TENSORFLOW_USE_ROCM + deepmd::tabulate_fusion_se_a_grad_gpu_rocm( + dy_dem_x, dy_dem, table, table_info, em_x, em, two_embed, dy, nloc, + nnei, last_layer_size); +#endif // TENSORFLOW_USE_ROCM + } else if (device == "CPU") { + deepmd::tabulate_fusion_se_a_grad_cpu(dy_dem_x, dy_dem, table, table_info, + em_x, em, two_embed, dy, nloc, nnei, + last_layer_size); + } + } + + private: + std::string device; +}; + template class TabulateFusionSeTOp : public OpKernel { public: @@ -725,6 +905,13 @@ class TabulateFusionSeRGradGradOp : public OpKernel { .Device(DEVICE_CPU) \ .TypeConstraint("T"), \ TabulateFusionSeAGradGradOp); \ + REGISTER_KERNEL_BUILDER( \ + Name("TabulateFusionSeAtten").Device(DEVICE_CPU).TypeConstraint("T"), \ + TabulateFusionSeAttenOp); \ + REGISTER_KERNEL_BUILDER(Name("TabulateFusionSeAttenGrad") \ + .Device(DEVICE_CPU) \ + .TypeConstraint("T"), \ + TabulateFusionSeAttenGradOp); \ REGISTER_KERNEL_BUILDER( \ Name("TabulateFusionSeT").Device(DEVICE_CPU).TypeConstraint("T"), \ TabulateFusionSeTOp); \ @@ -780,6 +967,16 @@ REGISTER_CPU(double); .TypeConstraint("T") \ .HostMemory("table_info"), \ TabulateFusionSeAGradGradOp); \ + REGISTER_KERNEL_BUILDER(Name("TabulateFusionSeAtten") \ + .Device(DEVICE_GPU) \ + .TypeConstraint("T") \ + .HostMemory("table_info"), \ + TabulateFusionSeAttenOp); \ + REGISTER_KERNEL_BUILDER(Name("TabulateFusionSeAttenGrad") \ + .Device(DEVICE_GPU) \ + .TypeConstraint("T") \ + .HostMemory("table_info"), \ + TabulateFusionSeAttenGradOp); \ REGISTER_KERNEL_BUILDER(Name("TabulateFusionSeT") \ .Device(DEVICE_GPU) \ .TypeConstraint("T") \ diff --git a/source/tests/test_data_large_batch.py b/source/tests/test_data_large_batch.py index 531c6de4f3..26b3a2957d 100644 --- a/source/tests/test_data_large_batch.py +++ b/source/tests/test_data_large_batch.py @@ -1,3 +1,4 @@ +import inspect import unittest import numpy as np @@ -41,6 +42,7 @@ class TestDataLargeBatch(tf.test.TestCase): def setUp(self): gen_data(mixed_type=True) + self.filename = __file__ def test_data_mixed_type(self): jfile = "water_se_atten_mixed_type.json" @@ -155,7 +157,7 @@ def test_data_mixed_type(self): t_box, t_mesh, inputs_dict, - suffix="se_atten", + suffix=self.filename + "-" + inspect.stack()[0][3], reuse=False, ) @@ -235,3 +237,395 @@ def test_data_mixed_type(self): np.testing.assert_almost_equal(e, refe, places) np.testing.assert_almost_equal(f, reff, places) np.testing.assert_almost_equal(v, refv, places) + sess.close() + + def test_stripped_data_mixed_type(self): + jfile = "water_se_atten_mixed_type.json" + jdata = j_loader(jfile) + + systems = j_must_have(jdata, "systems") + batch_size = 1 + test_size = 1 + rcut = j_must_have(jdata["model"]["descriptor"], "rcut") + type_map = j_must_have(jdata["model"], "type_map") + + data = DeepmdDataSystem(systems, batch_size, test_size, rcut, type_map=type_map) + 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) + + test_data = data.get_test() + numb_test = 1 + + jdata["model"]["descriptor"].pop("type", None) + jdata["model"]["descriptor"]["ntypes"] = 2 + jdata["model"]["descriptor"]["stripped_type_embedding"] = True + descrpt = DescrptSeAtten(**jdata["model"]["descriptor"], uniform_seed=True) + jdata["model"]["fitting_net"]["descrpt"] = descrpt + fitting = EnerFitting(**jdata["model"]["fitting_net"], uniform_seed=True) + typeebd_param = jdata["model"]["type_embedding"] + typeebd = TypeEmbedNet( + neuron=typeebd_param["neuron"], + resnet_dt=typeebd_param["resnet_dt"], + activation_function=None, + seed=typeebd_param["seed"], + uniform_seed=True, + padding=True, + ) + model = EnerModel(descrpt, fitting, typeebd) + + # 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"]], + "real_natoms_vec": [test_data["real_natoms_vec"]], + } + model._compute_input_stat(input_data, mixed_type=True) + model.descrpt.bias_atom_e = np.array([0.0, 0.0]) + + 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 + inputs_dict = {} + + model_pred = model.build( + t_coord, + t_type, + t_natoms, + t_box, + t_mesh, + inputs_dict, + suffix=self.filename + "-" + inspect.stack()[0][3], + reuse=False, + ) + + energy = model_pred["energy"] + force = model_pred["force"] + virial = model_pred["virial"] + atom_ener = model_pred["atom_ener"] + + feed_dict_test = { + t_energy: np.reshape(test_data["energy"][:numb_test], [-1]), + 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 = self.test_session().__enter__() + sess.run(tf.global_variables_initializer()) + [e, f, v] = sess.run([energy, force, virial], feed_dict=feed_dict_test) + # print(sess.run(model.type_embedding)) + # np.savetxt('tmp.out', sess.run(descrpt.dout, feed_dict = feed_dict_test), fmt='%.10e') + # # print(sess.run(model.atype_embed, feed_dict = feed_dict_test)) + # print(sess.run(fitting.inputs, feed_dict = feed_dict_test)) + # print(sess.run(fitting.outs, feed_dict = feed_dict_test)) + # print(sess.run(fitting.atype_embed, feed_dict = feed_dict_test)) + + e = e.reshape([-1]) + f = f.reshape([-1]) + v = v.reshape([-1]) + np.savetxt("e11.out", e.reshape([1, -1]), delimiter=",") + np.savetxt("f11.out", f.reshape([1, -1]), delimiter=",") + np.savetxt("v11.out", v.reshape([1, -1]), delimiter=",") + + refe = [6.125926357944699419e01] + reff = [ + 4.071033392194846855e-03, + 1.191078506555811808e-02, + 5.710038490045591959e-04, + 2.083813511902148086e-02, + 1.050404909007916256e-02, + -1.935131519230624082e-03, + -1.844253334357196135e-03, + 7.073208513688628192e-03, + -5.000418009101099666e-05, + 2.877594036828151017e-03, + -1.849276075329028823e-02, + -5.424378676202407318e-04, + -3.237425532982485255e-02, + 2.747768700765259881e-03, + 2.122946741188093227e-03, + 6.431746116137571252e-03, + -1.374305061680087571e-02, + -1.663770232507767613e-04, + ] + refv = [ + -5.699812217912397783e-02, + 8.904976757403395421e-03, + 2.306167440955633578e-03, + 8.904976757403414503e-03, + -3.502855693434053092e-02, + -6.596869271547717252e-04, + 2.306167440955633145e-03, + -6.596869271547715083e-04, + -1.602012510288682446e-04, + ] + + refe = np.reshape(refe, [-1]) + reff = np.reshape(reff, [-1]) + refv = np.reshape(refv, [-1]) + + places = 10 + np.testing.assert_almost_equal(e, refe, places) + np.testing.assert_almost_equal(f, reff, places) + np.testing.assert_almost_equal(v, refv, places) + + def test_compressible_data_mixed_type(self): + jfile = "water_se_atten_mixed_type.json" + jdata = j_loader(jfile) + + systems = j_must_have(jdata, "systems") + batch_size = 1 + test_size = 1 + rcut = j_must_have(jdata["model"]["descriptor"], "rcut") + type_map = j_must_have(jdata["model"], "type_map") + + data = DeepmdDataSystem(systems, batch_size, test_size, rcut, type_map=type_map) + 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) + + test_data = data.get_test() + numb_test = 1 + + jdata["model"]["descriptor"].pop("type", None) + jdata["model"]["descriptor"]["ntypes"] = 2 + jdata["model"]["descriptor"]["stripped_type_embedding"] = True + jdata["model"]["descriptor"]["attn_layer"] = 0 + descrpt = DescrptSeAtten(**jdata["model"]["descriptor"], uniform_seed=True) + jdata["model"]["fitting_net"]["descrpt"] = descrpt + fitting = EnerFitting(**jdata["model"]["fitting_net"], uniform_seed=True) + typeebd_param = jdata["model"]["type_embedding"] + typeebd = TypeEmbedNet( + neuron=typeebd_param["neuron"], + resnet_dt=typeebd_param["resnet_dt"], + activation_function=None, + seed=typeebd_param["seed"], + uniform_seed=True, + padding=True, + ) + model = EnerModel(descrpt, fitting, typeebd) + + # 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"]], + "real_natoms_vec": [test_data["real_natoms_vec"]], + } + model._compute_input_stat(input_data, mixed_type=True) + model.descrpt.bias_atom_e = np.array([0.0, 0.0]) + + 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 + inputs_dict = {} + + model_pred = model.build( + t_coord, + t_type, + t_natoms, + t_box, + t_mesh, + inputs_dict, + suffix=self.filename + "-" + inspect.stack()[0][3], + reuse=False, + ) + + energy = model_pred["energy"] + force = model_pred["force"] + virial = model_pred["virial"] + atom_ener = model_pred["atom_ener"] + + feed_dict_test = { + t_energy: np.reshape(test_data["energy"][:numb_test], [-1]), + 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 = self.test_session().__enter__() + sess.run(tf.global_variables_initializer()) + [e, f, v] = sess.run([energy, force, virial], feed_dict=feed_dict_test) + # print(sess.run(model.type_embedding)) + # np.savetxt('tmp.out', sess.run(descrpt.dout, feed_dict = feed_dict_test), fmt='%.10e') + # # print(sess.run(model.atype_embed, feed_dict = feed_dict_test)) + # print(sess.run(fitting.inputs, feed_dict = feed_dict_test)) + # print(sess.run(fitting.outs, feed_dict = feed_dict_test)) + # print(sess.run(fitting.atype_embed, feed_dict = feed_dict_test)) + + e = e.reshape([-1]) + f = f.reshape([-1]) + v = v.reshape([-1]) + np.savetxt("e.out", e.reshape([1, -1]), delimiter=",") + np.savetxt("f.out", f.reshape([1, -1]), delimiter=",") + np.savetxt("v.out", v.reshape([1, -1]), delimiter=",") + + refe = [4.978373241868134613e01] + reff = [ + 3.587688614359243466e00, + 3.202584939641652362e00, + 1.166711402014127957e-01, + 2.384342214774975321e00, + 3.542611694579458348e00, + -1.916097942322055603e-01, + -4.120123413353201869e00, + 1.474564563185293276e00, + 2.693540383300669847e-02, + 2.380464377433281431e00, + -4.807108079981841975e00, + -1.784915273650321821e-01, + -5.314498717923408222e00, + 1.495140750360528958e00, + 2.602480033292806638e-01, + 1.082126924709109872e00, + -4.907793867785092523e00, + -3.375322576646228034e-02, + ] + refv = [ + -1.760844499856655432e01, + 3.767507287595555532e00, + 4.505304110104397242e-01, + 3.767507287595555088e00, + -1.052518611764145362e01, + -2.174256785611231313e-01, + 4.505304110104397797e-01, + -2.174256785611231868e-01, + -2.462288791771098315e-02, + ] + + refe = np.reshape(refe, [-1]) + reff = np.reshape(reff, [-1]) + refv = np.reshape(refv, [-1]) + + places = 10 + np.testing.assert_almost_equal(e, refe, places) + np.testing.assert_almost_equal(f, reff, places) + np.testing.assert_almost_equal(v, refv, places) diff --git a/source/tests/test_descrpt_se_atten.py b/source/tests/test_descrpt_se_atten.py index cfa2cdaae5..f0b45e5dee 100644 --- a/source/tests/test_descrpt_se_atten.py +++ b/source/tests/test_descrpt_se_atten.py @@ -1,3 +1,4 @@ +import inspect import unittest import numpy as np @@ -33,6 +34,7 @@ class TestModel(tf.test.TestCase): def setUp(self): gen_data(nframes=2) + self.filename = __file__ def test_descriptor_two_sides(self): jfile = "water_se_atten.json" @@ -102,7 +104,13 @@ def test_descriptor_two_sides(self): is_training = tf.placeholder(tf.bool) t_fparam = None - type_embedding = typeebd.build(ntypes, suffix="_se_atten_type_des_ebd_2sdies") + type_embedding = typeebd.build( + ntypes, + suffix=self.filename + + "-" + + inspect.stack()[0][3] + + "_se_atten_type_des_ebd_2sdies", + ) dout = descrpt.build( t_coord, @@ -112,7 +120,10 @@ def test_descriptor_two_sides(self): t_mesh, {"type_embedding": type_embedding}, reuse=False, - suffix="_se_atten_type_des_2sides", + suffix=self.filename + + "-" + + inspect.stack()[0][3] + + "_se_atten_type_des_2sides", ) feed_dict_test = { @@ -133,6 +144,7 @@ def test_descriptor_two_sides(self): sess.run(tf.global_variables_initializer()) [model_dout] = sess.run([dout], feed_dict=feed_dict_test) model_dout = model_dout.reshape([-1]) + np.savetxt("two.out", model_dout.reshape([1, -1]), delimiter=",") ref_dout = [ 1.3503570575883254e-04, @@ -269,7 +281,13 @@ def test_descriptor_one_side(self): is_training = tf.placeholder(tf.bool) t_fparam = None - type_embedding = typeebd.build(ntypes, suffix="_se_atten_type_des_ebd_1side") + type_embedding = typeebd.build( + ntypes, + suffix=self.filename + + "-" + + inspect.stack()[0][3] + + "_se_atten_type_des_ebd_1side", + ) dout = descrpt.build( t_coord, @@ -279,7 +297,10 @@ def test_descriptor_one_side(self): t_mesh, {"type_embedding": type_embedding}, reuse=False, - suffix="_se_atten_type_des_1side", + suffix=self.filename + + "-" + + inspect.stack()[0][3] + + "_se_atten_type_des_1side", ) feed_dict_test = { @@ -300,6 +321,7 @@ def test_descriptor_one_side(self): sess.run(tf.global_variables_initializer()) [model_dout] = sess.run([dout], feed_dict=feed_dict_test) model_dout = model_dout.reshape([-1]) + np.savetxt("one.out", model_dout.reshape([1, -1]), delimiter=",") ref_dout = [ 8.9336098555659429e-05, @@ -366,3 +388,351 @@ def test_descriptor_one_side(self): places = 10 np.testing.assert_almost_equal(model_dout, ref_dout, places) + + def test_stripped_type_embedding_descriptor_two_sides(self): + jfile = "water_se_atten.json" + jdata = j_loader(jfile) + + 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 = 2 + test_size = 1 + stop_batch = j_must_have(jdata, "stop_batch") + rcut = j_must_have(jdata["model"]["descriptor"], "rcut") + sel = j_must_have(jdata["model"]["descriptor"], "sel") + ntypes = len(jdata["model"]["type_map"]) + + data = DataSystem(systems, set_pfx, batch_size, test_size, rcut, run_opt=None) + + test_data = data.get_test() + numb_test = 1 + + # set parameters + jdata["model"]["descriptor"]["neuron"] = [5, 5, 5] + jdata["model"]["descriptor"]["axis_neuron"] = 2 + typeebd_param = { + "neuron": [5], + "resnet_dt": False, + "seed": 1, + } + jdata["model"]["descriptor"]["stripped_type_embedding"] = True + + # init models + typeebd = TypeEmbedNet( + neuron=typeebd_param["neuron"], + activation_function=None, + resnet_dt=typeebd_param["resnet_dt"], + seed=typeebd_param["seed"], + uniform_seed=True, + padding=True, + ) + + jdata["model"]["descriptor"].pop("type", None) + jdata["model"]["descriptor"]["ntypes"] = ntypes + descrpt = DescrptSeAtten(**jdata["model"]["descriptor"], uniform_seed=True) + + # 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"]], + } + 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, [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 + + type_embedding = typeebd.build( + ntypes, suffix=self.filename + "-" + inspect.stack()[0][3] + ) + + dout = descrpt.build( + t_coord, + t_type, + t_natoms, + t_box, + t_mesh, + {"type_embedding": type_embedding}, + reuse=False, + suffix=self.filename + "-" + inspect.stack()[0][3], + ) + + 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 = self.test_session().__enter__() + sess.run(tf.global_variables_initializer()) + [model_dout] = sess.run([dout], feed_dict=feed_dict_test) + model_dout = model_dout.reshape([-1]) + np.savetxt("two1.out", model_dout.reshape([1, -1]), delimiter=",") + + ref_dout = [ + 6.383405295803018563e-06, + -4.970443805148654732e-05, + -4.970443805148654732e-05, + 3.915685888463831281e-04, + -4.475044754350870549e-05, + 3.524944903606048942e-04, + -3.724785876580149628e-05, + 2.922698210610697768e-04, + 1.253193390649937305e-04, + -9.866284622165712517e-04, + 2.746706080574298427e-06, + -2.936129326691491197e-05, + -2.936129326691491197e-05, + 3.361216906441433726e-04, + -2.672646322529656884e-05, + 3.054685398038972568e-04, + -1.987664884734486549e-05, + 2.233066336404597018e-04, + 7.321769925898204405e-05, + -8.355355708215854091e-04, + 2.577054933601801410e-06, + -2.596197919081208680e-05, + -2.596197919081208680e-05, + 2.634209681343336554e-04, + -2.437422597229959228e-05, + 2.471257526952611570e-04, + -1.804241933402407212e-05, + 1.824137020231941916e-04, + 6.580156956353395276e-05, + -6.669984436619770833e-04, + 7.430461358270195101e-06, + -4.911199551777768784e-05, + -4.911199551777768784e-05, + 3.246586591299433010e-04, + -4.545306022962886084e-05, + 3.004976436904362321e-04, + -3.943764170159576608e-05, + 2.607633588556588491e-04, + 1.265722360907320951e-04, + -8.368076661582606740e-04, + 1.006529590251595632e-05, + -5.947201651377735437e-05, + -5.947201651377735437e-05, + 3.519072130827992253e-04, + -5.475453207501003249e-05, + 3.240000254870917815e-04, + -5.001544285922251417e-05, + 2.958951751684798311e-04, + 1.541766955454939312e-04, + -9.123303972245934293e-04, + 5.302564262473944598e-06, + -4.087818220772074430e-05, + -4.087818220772074430e-05, + 3.170373643181760982e-04, + -3.805426060094593870e-05, + 2.950399422301473447e-04, + -3.093771463791340635e-05, + 2.394374838886431352e-04, + 1.045675931841061270e-04, + -8.106366082292457186e-04, + ] + + places = 10 + np.testing.assert_almost_equal(model_dout, ref_dout, places) + + def test_compressible_descriptor_two_sides(self): + jfile = "water_se_atten.json" + jdata = j_loader(jfile) + + 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 = 2 + test_size = 1 + stop_batch = j_must_have(jdata, "stop_batch") + rcut = j_must_have(jdata["model"]["descriptor"], "rcut") + sel = j_must_have(jdata["model"]["descriptor"], "sel") + ntypes = len(jdata["model"]["type_map"]) + + data = DataSystem(systems, set_pfx, batch_size, test_size, rcut, run_opt=None) + + test_data = data.get_test() + numb_test = 1 + + # set parameters + jdata["model"]["descriptor"]["neuron"] = [5, 5, 5] + jdata["model"]["descriptor"]["axis_neuron"] = 2 + jdata["model"]["descriptor"]["attn_layer"] = 0 + jdata["model"]["descriptor"]["stripped_type_embedding"] = True + typeebd_param = { + "neuron": [5], + "resnet_dt": False, + "seed": 1, + } + + # init models + typeebd = TypeEmbedNet( + neuron=typeebd_param["neuron"], + activation_function=None, + resnet_dt=typeebd_param["resnet_dt"], + seed=typeebd_param["seed"], + uniform_seed=True, + padding=True, + ) + + jdata["model"]["descriptor"].pop("type", None) + jdata["model"]["descriptor"]["ntypes"] = ntypes + descrpt = DescrptSeAtten(**jdata["model"]["descriptor"], uniform_seed=True) + + # 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"]], + } + 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, [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 + + type_embedding = typeebd.build( + ntypes, + suffix=self.filename + + "-" + + inspect.stack()[0][3] + + "_se_atten_type_des_ebd_2sdies", + ) + + dout = descrpt.build( + t_coord, + t_type, + t_natoms, + t_box, + t_mesh, + {"type_embedding": type_embedding}, + reuse=False, + suffix=self.filename + + "-" + + inspect.stack()[0][3] + + "_se_atten_type_des_2sides", + ) + + 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 = self.test_session().__enter__() + sess.run(tf.global_variables_initializer()) + [model_dout] = sess.run([dout], feed_dict=feed_dict_test) + model_dout = model_dout.reshape([-1]) + np.savetxt("two.out", model_dout.reshape([1, -1]), delimiter=",") + + ref_dout = [ + 1.141503478024449673e-02, + -3.178281748950724847e-05, + -3.178281748950724847e-05, + 9.254130621715512665e-08, + 9.355933444157461115e-03, + -2.610048686000205533e-05, + 4.941903103720379602e-03, + -1.380808118210685460e-05, + 3.126563600452313824e-02, + -8.708904919620347271e-05, + 9.450853881799477565e-03, + -2.980684015958201359e-05, + -2.980684015958201359e-05, + 1.108555048277460046e-07, + 7.781618028607465050e-03, + -2.477633432327578414e-05, + 4.121729175224964070e-03, + -1.322807700672256342e-05, + 2.594071094608556283e-02, + -8.191466287679302258e-05, + 8.687152275749851840e-03, + 7.169964582221861242e-04, + 7.169964582221861242e-04, + 5.920093158205708974e-05, + 6.980509354575046600e-03, + 5.760687878669166736e-04, + 3.535108314119884531e-03, + 2.917691604627718586e-04, + 2.237560827369438829e-02, + 1.845969886863443182e-03, + 1.117485257493494588e-02, + 9.236674604597058238e-04, + 9.236674604597058238e-04, + 7.636846801269864835e-05, + 8.961677239268169004e-03, + 7.406856311644359698e-04, + 4.531907674672685724e-03, + 3.746091075776196411e-04, + 2.878475290103474707e-02, + 2.378376253052099428e-03, + 1.227893341898089741e-02, + 1.013435353901762093e-03, + 1.013435353901762093e-03, + 8.366615519789353964e-05, + 9.853113810276110246e-03, + 8.131104560097670826e-04, + 4.984313719158160738e-03, + 4.113088418105881284e-04, + 3.165040520299260834e-02, + 2.611774507262543468e-03, + 1.073868692816030805e-02, + 8.884306507846149482e-04, + 8.884306507846149482e-04, + 7.351499010825891778e-05, + 8.615254842259736923e-03, + 7.126675675525927581e-04, + 4.359946655901846685e-03, + 3.606316239946267860e-04, + 2.762742703322493970e-02, + 2.285540100738304818e-03, + ] + + places = 10 + np.testing.assert_almost_equal(model_dout, ref_dout, places) diff --git a/source/tests/test_finetune_se_atten.py b/source/tests/test_finetune_se_atten.py index 18c89f0153..0631694027 100644 --- a/source/tests/test_finetune_se_atten.py +++ b/source/tests/test_finetune_se_atten.py @@ -55,11 +55,13 @@ def _subprocess_run(command): return popen.returncode -def _init_models(): +def _init_models(setup_model, i): data_file = str(tests_path / os.path.join("finetune", "data")) data_file_mixed_type = str(tests_path / os.path.join("finetune", "data_mixed_type")) pretrained_model = str(tests_path / "pretrained_model_se_atten.pb") finetuned_model = str(tests_path / "finetuned_model_se_atten.pb") + # ckpt_pretrain = str(tests_path / f"pretrain{i}.ckpt") + # ckpt_finetune = str(tests_path / f"finetune{i}.ckpt") finetuned_model_mixed_type = str( tests_path / "finetuned_model_se_atten_mixed_type.pb" ) @@ -74,8 +76,13 @@ def _init_models(): ) jdata_pre["training"]["training_data"]["systems"] = data_file jdata_pre["training"]["validation_data"]["systems"] = data_file + # jdata_pre["training"]["save_ckpt"] = ckpt_pretrain + setup_model(jdata_pre) + jdata_finetune["training"]["training_data"]["systems"] = data_file jdata_finetune["training"]["validation_data"]["systems"] = data_file + # jdata_finetune["training"]["save_ckpt"] = ckpt_finetune + setup_model(jdata_finetune) type_map_pre = jdata_pre["model"]["type_map"] type_map_finetune = jdata_finetune["model"]["type_map"] with open(INPUT_PRE, "w") as fp: @@ -140,17 +147,50 @@ def _init_models(): if not parse_version(tf.__version__) < parse_version("1.15"): - ( - INPUT_PRE, - INPUT_FINETUNE, - INPUT_FINETUNE_MIX, - PRE_MODEL, - FINETUNED_MODEL, - FINETUNED_MODEL_MIX, - PRE_MAP, - FINETUNED_MAP, - VALID_DATA, - ) = _init_models() + + def previous_se_atten(jdata): + jdata["model"]["descriptor"]["stripped_type_embedding"] = False + jdata["model"]["descriptor"]["attn_layer"] = 2 + + def stripped_model(jdata): + jdata["model"]["descriptor"]["stripped_type_embedding"] = True + jdata["model"]["descriptor"]["attn_layer"] = 2 + + def compressible_model(jdata): + jdata["model"]["descriptor"]["stripped_type_embedding"] = True + jdata["model"]["descriptor"]["attn_layer"] = 0 + + models = [previous_se_atten, stripped_model, compressible_model] + INPUT_PRES = [] + INPUT_FINETUNES = [] + INPUT_FINETUNE_MIXS = [] + PRE_MODELS = [] + FINETUNED_MODELS = [] + FINETUNED_MODEL_MIXS = [] + PRE_MAPS = [] + FINETUNED_MAPS = [] + VALID_DATAS = [] + for i, model in enumerate(models): + ( + INPUT_PRE, + INPUT_FINETUNE, + INPUT_FINETUNE_MIX, + PRE_MODEL, + FINETUNED_MODEL, + FINETUNED_MODEL_MIX, + PRE_MAP, + FINETUNED_MAP, + VALID_DATA, + ) = _init_models(model, i) + INPUT_PRES.append(INPUT_PRE) + INPUT_FINETUNES.append(INPUT_FINETUNE) + INPUT_FINETUNE_MIXS.append(INPUT_FINETUNE_MIX) + PRE_MODELS.append(PRE_MODEL) + FINETUNED_MODELS.append(FINETUNED_MODEL) + FINETUNED_MODEL_MIXS.append(FINETUNED_MODEL_MIX) + PRE_MAPS.append(PRE_MAP) + FINETUNED_MAPS.append(FINETUNED_MAP) + VALID_DATAS.append(VALID_DATA) @unittest.skipIf( @@ -158,96 +198,111 @@ def _init_models(): f"The current tf version {tf.__version__} is too low to run the new testing model.", ) class TestFinetuneSeAtten(unittest.TestCase): - @classmethod - def setUpClass(self): - self.valid_data = VALID_DATA - @classmethod def tearDownClass(self): - _file_delete(INPUT_PRE) - _file_delete(INPUT_FINETUNE) - _file_delete(INPUT_FINETUNE_MIX) - _file_delete(PRE_MODEL) - _file_delete(FINETUNED_MODEL) - _file_delete(FINETUNED_MODEL_MIX) - _file_delete("out.json") - _file_delete(str(tests_path / "checkpoint")) - _file_delete("model.ckpt.meta") - _file_delete("model.ckpt.index") - _file_delete("model.ckpt.data-00000-of-00001") - _file_delete("model.ckpt-0.meta") - _file_delete("model.ckpt-0.index") - _file_delete("model.ckpt-0.data-00000-of-00001") - _file_delete("model.ckpt-1.meta") - _file_delete("model.ckpt-1.index") - _file_delete("model.ckpt-1.data-00000-of-00001") - _file_delete("input_v2_compat.json") - _file_delete("lcurve.out") + for i in range(len(INPUT_PRES)): + _file_delete(INPUT_PRES[i]) + _file_delete(INPUT_FINETUNES[i]) + _file_delete(INPUT_FINETUNE_MIXS[i]) + _file_delete(PRE_MODELS[i]) + _file_delete(FINETUNED_MODELS[i]) + _file_delete(FINETUNED_MODEL_MIXS[i]) + _file_delete("out.json") + _file_delete("model.ckpt.meta") + _file_delete("model.ckpt.index") + _file_delete("model.ckpt.data-00000-of-00001") + _file_delete("model.ckpt-0.meta") + _file_delete("model.ckpt-0.index") + _file_delete("model.ckpt-0.data-00000-of-00001") + _file_delete("model.ckpt-1.meta") + _file_delete("model.ckpt-1.index") + _file_delete("model.ckpt-1.data-00000-of-00001") + _file_delete(str(tests_path / "checkpoint")) + _file_delete("input_v2_compat.json") + _file_delete("lcurve.out") def test_finetune_standard(self): - pretrained_bias = get_tensor_by_name(PRE_MODEL, "fitting_attr/t_bias_atom_e") - finetuned_bias = get_tensor_by_name( - FINETUNED_MODEL, "fitting_attr/t_bias_atom_e" - ) - sorter = np.argsort(PRE_MAP) - idx_type_map = sorter[np.searchsorted(PRE_MAP, FINETUNED_MAP, sorter=sorter)] - test_data = self.valid_data.get_test() - atom_nums = np.tile(np.bincount(test_data["type"][0])[idx_type_map], (4, 1)) - - dp = DeepPotential(PRE_MODEL) - energy = dp.eval(test_data["coord"], test_data["box"], test_data["type"][0])[0] - energy_diff = test_data["energy"] - energy - finetune_shift = finetuned_bias[idx_type_map] - pretrained_bias[idx_type_map] - ground_truth_shift = np.linalg.lstsq(atom_nums, energy_diff, rcond=None)[ - 0 - ].reshape(-1) - - dp_finetuned = DeepPotential(FINETUNED_MODEL) - energy_finetuned = dp_finetuned.eval( - test_data["coord"], test_data["box"], test_data["type"][0] - )[0] - energy_diff_finetuned = test_data["energy"] - energy_finetuned - finetune_results = np.linalg.lstsq( - atom_nums, energy_diff_finetuned, rcond=None - )[0].reshape(-1) - - # check values - np.testing.assert_almost_equal( - finetune_shift, ground_truth_shift, default_places - ) - np.testing.assert_almost_equal(finetune_results, 0.0, default_places) + for i in range(len(INPUT_PRES)): + self.valid_data = VALID_DATAS[i] + pretrained_bias = get_tensor_by_name( + PRE_MODELS[i], "fitting_attr/t_bias_atom_e" + ) + finetuned_bias = get_tensor_by_name( + FINETUNED_MODELS[i], "fitting_attr/t_bias_atom_e" + ) + sorter = np.argsort(PRE_MAPS[i]) + idx_type_map = sorter[ + np.searchsorted(PRE_MAPS[i], FINETUNED_MAPS[i], sorter=sorter) + ] + test_data = self.valid_data.get_test() + atom_nums = np.tile(np.bincount(test_data["type"][0])[idx_type_map], (4, 1)) + + dp = DeepPotential(PRE_MODELS[i]) + energy = dp.eval( + test_data["coord"], test_data["box"], test_data["type"][0] + )[0] + energy_diff = test_data["energy"] - energy + finetune_shift = ( + finetuned_bias[idx_type_map] - pretrained_bias[idx_type_map] + ) + ground_truth_shift = np.linalg.lstsq(atom_nums, energy_diff, rcond=None)[ + 0 + ].reshape(-1) + + dp_finetuned = DeepPotential(FINETUNED_MODELS[i]) + energy_finetuned = dp_finetuned.eval( + test_data["coord"], test_data["box"], test_data["type"][0] + )[0] + energy_diff_finetuned = test_data["energy"] - energy_finetuned + finetune_results = np.linalg.lstsq( + atom_nums, energy_diff_finetuned, rcond=None + )[0].reshape(-1) + + # check values + np.testing.assert_almost_equal( + finetune_shift, ground_truth_shift, default_places + ) + np.testing.assert_almost_equal(finetune_results, 0.0, default_places) def test_finetune_mixed_type(self): - pretrained_bias = get_tensor_by_name(PRE_MODEL, "fitting_attr/t_bias_atom_e") - finetuned_bias_mixed_type = get_tensor_by_name( - FINETUNED_MODEL_MIX, "fitting_attr/t_bias_atom_e" - ) - sorter = np.argsort(PRE_MAP) - idx_type_map = sorter[np.searchsorted(PRE_MAP, FINETUNED_MAP, sorter=sorter)] - test_data = self.valid_data.get_test() - atom_nums = np.tile(np.bincount(test_data["type"][0])[idx_type_map], (4, 1)) - - dp = DeepPotential(PRE_MODEL) - energy = dp.eval(test_data["coord"], test_data["box"], test_data["type"][0])[0] - energy_diff = test_data["energy"] - energy - finetune_shift = ( - finetuned_bias_mixed_type[idx_type_map] - pretrained_bias[idx_type_map] - ) - ground_truth_shift = np.linalg.lstsq(atom_nums, energy_diff, rcond=None)[ - 0 - ].reshape(-1) - - dp_finetuned_mixed_type = DeepPotential(FINETUNED_MODEL_MIX) - energy_finetuned = dp_finetuned_mixed_type.eval( - test_data["coord"], test_data["box"], test_data["type"][0] - )[0] - energy_diff_finetuned = test_data["energy"] - energy_finetuned - finetune_results = np.linalg.lstsq( - atom_nums, energy_diff_finetuned, rcond=None - )[0].reshape(-1) - - # check values - np.testing.assert_almost_equal( - finetune_shift, ground_truth_shift, default_places - ) - np.testing.assert_almost_equal(finetune_results, 0.0, default_places) + for i in range(len(INPUT_PRES)): + self.valid_data = VALID_DATAS[i] + pretrained_bias = get_tensor_by_name( + PRE_MODELS[i], "fitting_attr/t_bias_atom_e" + ) + finetuned_bias_mixed_type = get_tensor_by_name( + FINETUNED_MODEL_MIXS[i], "fitting_attr/t_bias_atom_e" + ) + sorter = np.argsort(PRE_MAPS[i]) + idx_type_map = sorter[ + np.searchsorted(PRE_MAPS[i], FINETUNED_MAPS[i], sorter=sorter) + ] + test_data = self.valid_data.get_test() + atom_nums = np.tile(np.bincount(test_data["type"][0])[idx_type_map], (4, 1)) + + dp = DeepPotential(PRE_MODELS[i]) + energy = dp.eval( + test_data["coord"], test_data["box"], test_data["type"][0] + )[0] + energy_diff = test_data["energy"] - energy + finetune_shift = ( + finetuned_bias_mixed_type[idx_type_map] - pretrained_bias[idx_type_map] + ) + ground_truth_shift = np.linalg.lstsq(atom_nums, energy_diff, rcond=None)[ + 0 + ].reshape(-1) + + dp_finetuned_mixed_type = DeepPotential(FINETUNED_MODEL_MIXS[i]) + energy_finetuned = dp_finetuned_mixed_type.eval( + test_data["coord"], test_data["box"], test_data["type"][0] + )[0] + energy_diff_finetuned = test_data["energy"] - energy_finetuned + finetune_results = np.linalg.lstsq( + atom_nums, energy_diff_finetuned, rcond=None + )[0].reshape(-1) + + # check values + np.testing.assert_almost_equal( + finetune_shift, ground_truth_shift, default_places + ) + np.testing.assert_almost_equal(finetune_results, 0.0, default_places) diff --git a/source/tests/test_init_frz_model_se_atten.py b/source/tests/test_init_frz_model_se_atten.py index 86a4fd86b9..c623ec5017 100644 --- a/source/tests/test_init_frz_model_se_atten.py +++ b/source/tests/test_init_frz_model_se_atten.py @@ -43,10 +43,10 @@ def _file_delete(file): os.remove(file) -def _init_models(): +def _init_models(model_setup, i): 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") + frozen_model = str(tests_path / f"init_frz_se_atten{i}.pb") + ckpt = str(tests_path / f"init_frz_se_atten{i}.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") @@ -56,6 +56,7 @@ def _init_models(): jdata["training"]["save_ckpt"] = ckpt jdata["model"]["descriptor"]["type"] = "se_atten" jdata["model"]["descriptor"]["sel"] = 120 + model_setup(jdata) with open(INPUT, "w") as fp: json.dump(jdata, fp, indent=4) ret = run_dp("dp train " + INPUT) @@ -63,7 +64,7 @@ def _init_models(): ret = run_dp("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 = update_deepmd_input(jdata, warning=True, dump=f"input_v2_compat{i}.json") jdata = normalize(jdata) model_ckpt = DPTrainer(jdata, run_opt=run_opt_ckpt) model_frz = DPTrainer(jdata, run_opt=run_opt_frz) @@ -131,15 +132,44 @@ def _init_models(): if not parse_version(tf.__version__) < parse_version("1.15"): - ( - INPUT, - CKPT, - FROZEN_MODEL, - CKPT_TRAINER, - FRZ_TRAINER, - VALID_DATA, - STOP_BATCH, - ) = _init_models() + + def previous_se_atten(jdata): + jdata["model"]["descriptor"]["stripped_type_embedding"] = False + jdata["model"]["descriptor"]["attn_layer"] = 2 + + def stripped_model(jdata): + jdata["model"]["descriptor"]["stripped_type_embedding"] = True + jdata["model"]["descriptor"]["attn_layer"] = 2 + + def compressible_model(jdata): + jdata["model"]["descriptor"]["stripped_type_embedding"] = True + jdata["model"]["descriptor"]["attn_layer"] = 0 + + models = [previous_se_atten, stripped_model, compressible_model] + INPUTS = [] + CKPTS = [] + FROZEN_MODELS = [] + CKPT_TRAINERS = [] + FRZ_TRAINERS = [] + VALID_DATAS = [] + STOP_BATCHS = [] + for i, model in enumerate(models): + ( + INPUT, + CKPT, + FROZEN_MODEL, + CKPT_TRAINER, + FRZ_TRAINER, + VALID_DATA, + STOP_BATCH, + ) = _init_models(model, i) + INPUTS.append(INPUT) + CKPTS.append(CKPT) + FROZEN_MODELS.append(FROZEN_MODEL) + CKPT_TRAINERS.append(CKPT_TRAINER) + FRZ_TRAINERS.append(FRZ_TRAINER) + VALID_DATAS.append(VALID_DATA) + STOP_BATCHS.append(STOP_BATCH) @unittest.skipIf( @@ -149,54 +179,63 @@ def _init_models(): class TestInitFrzModelAtten(unittest.TestCase): @classmethod def setUpClass(cls): - cls.dp_ckpt = CKPT_TRAINER - cls.dp_frz = FRZ_TRAINER - cls.valid_data = VALID_DATA - cls.stop_batch = STOP_BATCH + cls.dp_ckpts = CKPT_TRAINERS + cls.dp_frzs = FRZ_TRAINERS + cls.valid_datas = VALID_DATAS + cls.stop_batchs = STOP_BATCHS @classmethod def tearDownClass(cls): - _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") + for i in range(len(cls.dp_ckpts)): + _file_delete(INPUTS[i]) + _file_delete(FROZEN_MODELS[i]) + _file_delete("out.json") + _file_delete(str(tests_path / "checkpoint")) + _file_delete(CKPT[i] + ".meta") + _file_delete(CKPT[i] + ".index") + _file_delete(CKPT[i] + ".data-00000-of-00001") + _file_delete(CKPT[i] + "-0.meta") + _file_delete(CKPT[i] + "-0.index") + _file_delete(CKPT[i] + "-0.data-00000-of-00001") + _file_delete(CKPT[i] + "-1.meta") + _file_delete(CKPT[i] + "-1.index") + _file_delete(CKPT[i] + "-1.data-00000-of-00001") + _file_delete(f"input_v2_compat{i}.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 - ) + for i in range(len(self.dp_ckpts)): + self.dp_ckpt = CKPT_TRAINERS[i] + self.dp_frz = FRZ_TRAINERS[i] + self.valid_data = VALID_DATAS[i] + self.stop_batch = STOP_BATCHS[i] + + 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_model_compression_se_atten.py b/source/tests/test_model_compression_se_atten.py new file mode 100644 index 0000000000..f2d4e5edad --- /dev/null +++ b/source/tests/test_model_compression_se_atten.py @@ -0,0 +1,666 @@ +import json +import os +import subprocess as sp +import unittest + +import numpy as np + +# from deepmd.entrypoints.compress import compress +from common import ( + j_loader, + run_dp, + tests_path, +) +from packaging.version import parse as parse_version + +from deepmd.env import ( + tf, +) +from deepmd.infer import ( + DeepPot, +) + + +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 + + +@unittest.skipIf( + parse_version(tf.__version__) < parse_version("2"), + f"The current tf version {tf.__version__} is too low to run the new testing model.", +) +def _init_models(): + data_file = str(tests_path / os.path.join("model_compression", "data")) + inputs, frozen_models, compressed_models = [], [], [] + # 4 tests: + # - type embedding FP64, se_atten FP64 + # - type embedding FP64, se_atten FP32 + # - type embedding FP32, se_atten FP64 + # - type embedding FP32, se_atten FP32 + tests = [ + {"se_atten precision": "float64", "type embedding precision": "float64"}, + {"se_atten precision": "float64", "type embedding precision": "float32"}, + {"se_atten precision": "float32", "type embedding precision": "float64"}, + {"se_atten precision": "float32", "type embedding precision": "float32"}, + ] + for i in range(4): + INPUT = str(tests_path / f"input{i}.json") + frozen_model = str(tests_path / f"dp-original-se-atten{i}.pb") + compressed_model = str(tests_path / f"dp-compressed-se-atten{i}.pb") + jdata = j_loader( + str(tests_path / os.path.join("model_compression", "input.json")) + ) + jdata["model"]["descriptor"] = {} + jdata["model"]["descriptor"]["type"] = "se_atten" + jdata["model"]["descriptor"]["precision"] = tests[i]["se_atten precision"] + jdata["model"]["descriptor"]["stripped_type_embedding"] = True + jdata["model"]["descriptor"]["sel"] = 120 + jdata["model"]["descriptor"]["attn_layer"] = 0 + jdata["model"]["type_embedding"] = {} + jdata["model"]["type_embedding"]["precision"] = tests[i][ + "type embedding precision" + ] + jdata["training"]["training_data"]["systems"] = data_file + jdata["training"]["validation_data"]["systems"] = data_file + with open(INPUT, "w") as fp: + json.dump(jdata, fp, indent=4) + + ret = run_dp("dp train " + INPUT) + np.testing.assert_equal(ret, 0, "DP train failed!") + ret = run_dp("dp freeze -o " + frozen_model) + np.testing.assert_equal(ret, 0, "DP freeze failed!") + ret = run_dp("dp compress " + " -i " + frozen_model + " -o " + compressed_model) + np.testing.assert_equal(ret, 0, "DP model compression failed!") + + inputs.append(INPUT) + frozen_models.append(frozen_model) + compressed_models.append(compressed_model) + + return inputs, frozen_models, compressed_models + + +INPUTS, FROZEN_MODELS, COMPRESSED_MODELS = _init_models() + + +def _get_default_places(nth_test): + return 10 if nth_test == 0 else 3 + + +@unittest.skipIf( + parse_version(tf.__version__) < parse_version("2"), + f"The current tf version {tf.__version__} is too low to run the new testing model.", +) +class TestDeepPotAPBC(unittest.TestCase): + @classmethod + def setUpClass(self): + self.dp_originals = [DeepPot(FROZEN_MODELS[i]) for i in range(4)] + self.dp_compresseds = [DeepPot(COMPRESSED_MODELS[i]) for i in range(4)] + self.coords = np.array( + [ + 12.83, + 2.56, + 2.18, + 12.09, + 2.87, + 2.74, + 00.25, + 3.32, + 1.68, + 3.36, + 3.00, + 1.81, + 3.51, + 2.51, + 2.60, + 4.27, + 3.22, + 1.56, + ] + ) + self.atype = [0, 1, 1, 0, 1, 1] + self.box = np.array([13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0]) + + def test_attrs(self): + for i in range(4): + dp_original = self.dp_originals[i] + dp_compressed = self.dp_compresseds[i] + default_places = _get_default_places(i) + + self.assertEqual(dp_original.get_ntypes(), 2) + self.assertAlmostEqual(dp_original.get_rcut(), 6.0, places=default_places) + self.assertEqual(dp_original.get_type_map(), ["O", "H"]) + self.assertEqual(dp_original.get_dim_fparam(), 0) + self.assertEqual(dp_original.get_dim_aparam(), 0) + + self.assertEqual(dp_compressed.get_ntypes(), 2) + self.assertAlmostEqual(dp_compressed.get_rcut(), 6.0, places=default_places) + self.assertEqual(dp_compressed.get_type_map(), ["O", "H"]) + self.assertEqual(dp_compressed.get_dim_fparam(), 0) + self.assertEqual(dp_compressed.get_dim_aparam(), 0) + + def test_1frame(self): + for i in range(4): + dp_original = self.dp_originals[i] + dp_compressed = self.dp_compresseds[i] + default_places = _get_default_places(i) + + ee0, ff0, vv0 = dp_original.eval( + self.coords, self.box, self.atype, atomic=False + ) + ee1, ff1, vv1 = dp_compressed.eval( + self.coords, self.box, self.atype, atomic=False + ) + # check shape of the returns + nframes = 1 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + # check values + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + def test_1frame_atm(self): + for i in range(4): + dp_original = self.dp_originals[i] + dp_compressed = self.dp_compresseds[i] + default_places = _get_default_places(i) + + ee0, ff0, vv0, ae0, av0 = dp_original.eval( + self.coords, self.box, self.atype, atomic=True + ) + ee1, ff1, vv1, ae1, av1 = dp_compressed.eval( + self.coords, self.box, self.atype, atomic=True + ) + # check shape of the returns + nframes = 1 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ae0.shape, (nframes, natoms, 1)) + self.assertEqual(av0.shape, (nframes, natoms, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + self.assertEqual(ae1.shape, (nframes, natoms, 1)) + self.assertEqual(av1.shape, (nframes, natoms, 9)) + # check values + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + def test_2frame_atm(self): + for i in range(4): + dp_original = self.dp_originals[i] + dp_compressed = self.dp_compresseds[i] + default_places = _get_default_places(i) + + coords2 = np.concatenate((self.coords, self.coords)) + box2 = np.concatenate((self.box, self.box)) + ee0, ff0, vv0, ae0, av0 = dp_original.eval( + coords2, box2, self.atype, atomic=True + ) + ee1, ff1, vv1, ae1, av1 = dp_compressed.eval( + coords2, box2, self.atype, atomic=True + ) + # check shape of the returns + nframes = 2 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ae0.shape, (nframes, natoms, 1)) + self.assertEqual(av0.shape, (nframes, natoms, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + self.assertEqual(ae1.shape, (nframes, natoms, 1)) + self.assertEqual(av1.shape, (nframes, natoms, 9)) + + # check values + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + +@unittest.skipIf( + parse_version(tf.__version__) < parse_version("2"), + f"The current tf version {tf.__version__} is too low to run the new testing model.", +) +class TestDeepPotANoPBC(unittest.TestCase): + @classmethod + def setUpClass(self): + self.dp_originals = [DeepPot(FROZEN_MODELS[i]) for i in range(4)] + self.dp_compresseds = [DeepPot(COMPRESSED_MODELS[i]) for i in range(4)] + self.coords = np.array( + [ + 12.83, + 2.56, + 2.18, + 12.09, + 2.87, + 2.74, + 00.25, + 3.32, + 1.68, + 3.36, + 3.00, + 1.81, + 3.51, + 2.51, + 2.60, + 4.27, + 3.22, + 1.56, + ] + ) + self.atype = [0, 1, 1, 0, 1, 1] + self.box = None + + def test_1frame(self): + for i in range(4): + dp_original = self.dp_originals[i] + dp_compressed = self.dp_compresseds[i] + default_places = _get_default_places(i) + + ee0, ff0, vv0 = dp_original.eval( + self.coords, self.box, self.atype, atomic=False + ) + ee1, ff1, vv1 = dp_compressed.eval( + self.coords, self.box, self.atype, atomic=False + ) + # check shape of the returns + nframes = 1 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + # check values + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + def test_1frame_atm(self): + for i in range(4): + dp_original = self.dp_originals[i] + dp_compressed = self.dp_compresseds[i] + default_places = _get_default_places(i) + + ee0, ff0, vv0, ae0, av0 = dp_original.eval( + self.coords, self.box, self.atype, atomic=True + ) + ee1, ff1, vv1, ae1, av1 = dp_compressed.eval( + self.coords, self.box, self.atype, atomic=True + ) + # check shape of the returns + nframes = 1 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ae0.shape, (nframes, natoms, 1)) + self.assertEqual(av0.shape, (nframes, natoms, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + self.assertEqual(ae1.shape, (nframes, natoms, 1)) + self.assertEqual(av1.shape, (nframes, natoms, 9)) + # check values + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + def test_2frame_atm(self): + for i in range(4): + dp_original = self.dp_originals[i] + dp_compressed = self.dp_compresseds[i] + default_places = _get_default_places(i) + + coords2 = np.concatenate((self.coords, self.coords)) + ee0, ff0, vv0, ae0, av0 = dp_original.eval( + coords2, self.box, self.atype, atomic=True + ) + ee1, ff1, vv1, ae1, av1 = dp_compressed.eval( + coords2, self.box, self.atype, atomic=True + ) + # check shape of the returns + nframes = 2 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ae0.shape, (nframes, natoms, 1)) + self.assertEqual(av0.shape, (nframes, natoms, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + self.assertEqual(ae1.shape, (nframes, natoms, 1)) + self.assertEqual(av1.shape, (nframes, natoms, 9)) + + # check values + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + +@unittest.skipIf( + parse_version(tf.__version__) < parse_version("2"), + f"The current tf version {tf.__version__} is too low to run the new testing model.", +) +class TestDeepPotALargeBoxNoPBC(unittest.TestCase): + @classmethod + def setUpClass(self): + self.dp_originals = [DeepPot(FROZEN_MODELS[i]) for i in range(4)] + self.dp_compresseds = [DeepPot(COMPRESSED_MODELS[i]) for i in range(4)] + self.coords = np.array( + [ + 12.83, + 2.56, + 2.18, + 12.09, + 2.87, + 2.74, + 00.25, + 3.32, + 1.68, + 3.36, + 3.00, + 1.81, + 3.51, + 2.51, + 2.60, + 4.27, + 3.22, + 1.56, + ] + ) + self.atype = [0, 1, 1, 0, 1, 1] + self.box = np.array([19.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0]) + + def test_1frame(self): + for i in range(4): + dp_original = self.dp_originals[i] + dp_compressed = self.dp_compresseds[i] + default_places = _get_default_places(i) + + ee0, ff0, vv0 = dp_original.eval( + self.coords, self.box, self.atype, atomic=False + ) + ee1, ff1, vv1 = dp_compressed.eval( + self.coords, self.box, self.atype, atomic=False + ) + # check shape of the returns + nframes = 1 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + # check values + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + def test_1frame_atm(self): + for i in range(4): + dp_original = self.dp_originals[i] + dp_compressed = self.dp_compresseds[i] + default_places = _get_default_places(i) + + ee0, ff0, vv0, ae0, av0 = dp_original.eval( + self.coords, self.box, self.atype, atomic=True + ) + ee1, ff1, vv1, ae1, av1 = dp_compressed.eval( + self.coords, self.box, self.atype, atomic=True + ) + # check shape of the returns + nframes = 1 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ae0.shape, (nframes, natoms, 1)) + self.assertEqual(av0.shape, (nframes, natoms, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + self.assertEqual(ae1.shape, (nframes, natoms, 1)) + self.assertEqual(av1.shape, (nframes, natoms, 9)) + # check values + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + def test_ase(self): + for i in range(4): + default_places = _get_default_places(i) + from ase import ( + Atoms, + ) + + from deepmd.calculator import ( + DP, + ) + + water0 = Atoms( + "OHHOHH", + positions=self.coords.reshape((-1, 3)), + cell=self.box.reshape((3, 3)), + calculator=DP(FROZEN_MODELS[i]), + ) + water1 = Atoms( + "OHHOHH", + positions=self.coords.reshape((-1, 3)), + cell=self.box.reshape((3, 3)), + calculator=DP(COMPRESSED_MODELS[i]), + ) + ee0 = water0.get_potential_energy() + ff0 = water0.get_forces() + ee1 = water1.get_potential_energy() + ff1 = water1.get_forces() + nframes = 1 + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + + +@unittest.skipIf( + parse_version(tf.__version__) < parse_version("2"), + f"The current tf version {tf.__version__} is too low to run the new testing model.", +) +class TestDeepPotAPBCExcludeTypes(unittest.TestCase): + @classmethod + def setUpClass(self): + self.dp_originals = [DeepPot(FROZEN_MODELS[i]) for i in range(4)] + self.dp_compresseds = [DeepPot(COMPRESSED_MODELS[i]) for i in range(4)] + self.coords = np.array( + [ + 12.83, + 2.56, + 2.18, + 12.09, + 2.87, + 2.74, + 00.25, + 3.32, + 1.68, + 3.36, + 3.00, + 1.81, + 3.51, + 2.51, + 2.60, + 4.27, + 3.22, + 1.56, + ] + ) + self.atype = [0, 1, 1, 0, 1, 1] + self.box = np.array([13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0]) + + @classmethod + def tearDownClass(self): + for i in range(4): + _file_delete(INPUTS[i]) + _file_delete(FROZEN_MODELS[i]) + _file_delete(COMPRESSED_MODELS[i]) + _file_delete("out.json") + _file_delete("compress.json") + _file_delete("checkpoint") + _file_delete("model.ckpt.meta") + _file_delete("model.ckpt.index") + _file_delete("model.ckpt.data-00000-of-00001") + _file_delete("model.ckpt-100.meta") + _file_delete("model.ckpt-100.index") + _file_delete("model.ckpt-100.data-00000-of-00001") + _file_delete("model-compression/checkpoint") + _file_delete("model-compression/model.ckpt.meta") + _file_delete("model-compression/model.ckpt.index") + _file_delete("model-compression/model.ckpt.data-00000-of-00001") + _file_delete("model-compression") + _file_delete("input_v2_compat.json") + _file_delete("lcurve.out") + + def test_attrs(self): + for i in range(4): + dp_original = self.dp_originals[i] + dp_compressed = self.dp_compresseds[i] + default_places = _get_default_places(i) + + self.assertEqual(dp_original.get_ntypes(), 2) + self.assertAlmostEqual(dp_original.get_rcut(), 6.0, places=default_places) + self.assertEqual(dp_original.get_type_map(), ["O", "H"]) + self.assertEqual(dp_original.get_dim_fparam(), 0) + self.assertEqual(dp_original.get_dim_aparam(), 0) + + self.assertEqual(dp_compressed.get_ntypes(), 2) + self.assertAlmostEqual(dp_compressed.get_rcut(), 6.0, places=default_places) + self.assertEqual(dp_compressed.get_type_map(), ["O", "H"]) + self.assertEqual(dp_compressed.get_dim_fparam(), 0) + self.assertEqual(dp_compressed.get_dim_aparam(), 0) + + def test_1frame(self): + for i in range(4): + dp_original = self.dp_originals[i] + dp_compressed = self.dp_compresseds[i] + default_places = _get_default_places(i) + + ee0, ff0, vv0 = dp_original.eval( + self.coords, self.box, self.atype, atomic=False + ) + ee1, ff1, vv1 = dp_compressed.eval( + self.coords, self.box, self.atype, atomic=False + ) + # check shape of the returns + nframes = 1 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + # check values + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + def test_1frame_atm(self): + for i in range(4): + dp_original = self.dp_originals[i] + dp_compressed = self.dp_compresseds[i] + default_places = _get_default_places(i) + + ee0, ff0, vv0, ae0, av0 = dp_original.eval( + self.coords, self.box, self.atype, atomic=True + ) + ee1, ff1, vv1, ae1, av1 = dp_compressed.eval( + self.coords, self.box, self.atype, atomic=True + ) + # check shape of the returns + nframes = 1 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ae0.shape, (nframes, natoms, 1)) + self.assertEqual(av0.shape, (nframes, natoms, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + self.assertEqual(ae1.shape, (nframes, natoms, 1)) + self.assertEqual(av1.shape, (nframes, natoms, 9)) + # check values + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + def test_2frame_atm(self): + for i in range(4): + dp_original = self.dp_originals[i] + dp_compressed = self.dp_compresseds[i] + default_places = _get_default_places(i) + + coords2 = np.concatenate((self.coords, self.coords)) + box2 = np.concatenate((self.box, self.box)) + ee0, ff0, vv0, ae0, av0 = dp_original.eval( + coords2, box2, self.atype, atomic=True + ) + ee1, ff1, vv1, ae1, av1 = dp_compressed.eval( + coords2, box2, self.atype, atomic=True + ) + # check shape of the returns + nframes = 2 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ae0.shape, (nframes, natoms, 1)) + self.assertEqual(av0.shape, (nframes, natoms, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + self.assertEqual(ae1.shape, (nframes, natoms, 1)) + self.assertEqual(av1.shape, (nframes, natoms, 9)) + + # check values + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) diff --git a/source/tests/test_model_se_atten.py b/source/tests/test_model_se_atten.py index b473c43c10..81a11a3900 100644 --- a/source/tests/test_model_se_atten.py +++ b/source/tests/test_model_se_atten.py @@ -1,3 +1,4 @@ +import inspect import unittest import numpy as np @@ -39,6 +40,7 @@ class TestModel(tf.test.TestCase): def setUp(self): gen_data() + self.filename = __file__ def test_model(self): jfile = "water_se_atten.json" @@ -108,7 +110,7 @@ def test_model(self): t_box, t_mesh, inputs_dict, - suffix="se_atten", + suffix="test_model_se_atten", reuse=False, ) energy = model_pred["energy"] @@ -232,7 +234,7 @@ def test_exclude_types(self): padding=True, ) type_embedding = typeebd.build( - ntypes, + ntypes, suffix=self.filename + "-" + inspect.stack()[0][3] ) dout = descrpt.build( t_coord, @@ -262,3 +264,464 @@ def test_exclude_types(self): np.testing.assert_almost_equal(des[:, 0:2], 0.0, 10) with self.assertRaises(AssertionError): np.testing.assert_almost_equal(des[:, 2:6], 0.0, 10) + + def test_compressible_model(self): + jfile = "water_se_atten.json" + jdata = j_loader(jfile) + + 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 + + jdata["model"]["descriptor"].pop("type", None) + jdata["model"]["descriptor"]["ntypes"] = 2 + jdata["model"]["descriptor"]["stripped_type_embedding"] = True + jdata["model"]["descriptor"]["attn_layer"] = 0 + descrpt = DescrptSeAtten(**jdata["model"]["descriptor"], uniform_seed=True) + jdata["model"]["fitting_net"]["descrpt"] = descrpt + fitting = EnerFitting(**jdata["model"]["fitting_net"], uniform_seed=True) + typeebd_param = jdata["model"]["type_embedding"] + typeebd = TypeEmbedNet( + neuron=typeebd_param["neuron"], + activation_function=None, + resnet_dt=typeebd_param["resnet_dt"], + seed=typeebd_param["seed"], + uniform_seed=True, + padding=True, + ) + model = EnerModel(descrpt, fitting, typeebd) + + # 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 + inputs_dict = {} + + model_pred = model.build( + t_coord, + t_type, + t_natoms, + t_box, + t_mesh, + inputs_dict, + suffix="test_model_se_atten_model_compressible", + 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 = self.test_session().__enter__() + sess.run(tf.global_variables_initializer()) + [e, f, v] = sess.run([energy, force, virial], feed_dict=feed_dict_test) + # print(sess.run(model.type_embedding)) + # np.savetxt('tmp.out', sess.run(descrpt.dout, feed_dict = feed_dict_test), fmt='%.10e') + # # print(sess.run(model.atype_embed, feed_dict = feed_dict_test)) + # print(sess.run(fitting.inputs, feed_dict = feed_dict_test)) + # print(sess.run(fitting.outs, feed_dict = feed_dict_test)) + # print(sess.run(fitting.atype_embed, feed_dict = feed_dict_test)) + + e = e.reshape([-1]) + f = f.reshape([-1]) + v = v.reshape([-1]) + np.savetxt("e.out", e.reshape([1, -1]), delimiter=",") + np.savetxt("f.out", f.reshape([1, -1]), delimiter=",") + np.savetxt("v.out", v.reshape([1, -1]), delimiter=",") + + refe = [4.978373241868134613e01] + reff = [ + 3.587688614359243466e00, + 3.202584939641652362e00, + 1.166711402014127957e-01, + 2.384342214774975321e00, + 3.542611694579458348e00, + -1.916097942322055603e-01, + -4.120123413353201869e00, + 1.474564563185293276e00, + 2.693540383300669847e-02, + 2.380464377433281431e00, + -4.807108079981841975e00, + -1.784915273650321821e-01, + -5.314498717923408222e00, + 1.495140750360528958e00, + 2.602480033292806638e-01, + 1.082126924709109872e00, + -4.907793867785092523e00, + -3.375322576646228034e-02, + ] + refv = [ + -1.760844499856656498e01, + 3.767507287595556420e00, + 4.505304110104396687e-01, + 3.767507287595555532e00, + -1.052518611764145540e01, + -2.174256785611232701e-01, + 4.505304110104397797e-01, + -2.174256785611233256e-01, + -2.462288791771097621e-02, + ] + + refe = np.reshape(refe, [-1]) + reff = np.reshape(reff, [-1]) + refv = np.reshape(refv, [-1]) + + places = 10 + np.testing.assert_almost_equal(e, refe, places) + np.testing.assert_almost_equal(f, reff, places) + np.testing.assert_almost_equal(v, refv, places) + + def test_compressible_exclude_types(self): + """In this test, we make type 0 has no interaction with type 0 and type 1, + so the descriptor should be zero for type 0 atoms. + """ + jfile = "water_se_atten.json" + jdata = j_loader(jfile) + + 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["model"]["descriptor"], "rcut") + ntypes = 2 + + data = DataSystem(systems, set_pfx, batch_size, test_size, rcut, run_opt=None) + + test_data = data.get_test() + numb_test = 1 + + # set parameters + jdata["model"]["descriptor"]["exclude_types"] = [[0, 0], [0, 1]] + + t_prop_c = tf.placeholder(tf.float32, [5], name="t_prop_c") + 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, [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) + + # successful + descrpt = DescrptSeAtten(ntypes=ntypes, **jdata["model"]["descriptor"]) + typeebd_param = jdata["model"]["type_embedding"] + jdata["model"]["descriptor"]["stripped_type_embedding"] = True + jdata["model"]["descriptor"]["attn_layer"] = 0 + typeebd = TypeEmbedNet( + neuron=typeebd_param["neuron"], + activation_function=None, + resnet_dt=typeebd_param["resnet_dt"], + seed=typeebd_param["seed"], + uniform_seed=True, + padding=True, + ) + type_embedding = typeebd.build( + ntypes, suffix=self.filename + "-" + inspect.stack()[0][3] + ) + dout = descrpt.build( + t_coord, + t_type, + t_natoms, + t_box, + t_mesh, + {"type_embedding": type_embedding}, + reuse=False, + suffix="_se_atten_compressible_exclude_types", + ) + + feed_dict_test1 = { + 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, + } + + with self.test_session() as sess: + sess.run(tf.global_variables_initializer()) + [des] = sess.run([dout], feed_dict=feed_dict_test1) + + np.testing.assert_almost_equal(des[:, 0:2], 0.0, 10) + with self.assertRaises(AssertionError): + np.testing.assert_almost_equal(des[:, 2:6], 0.0, 10) + + def test_stripped_type_embedding_model(self): + jfile = "water_se_atten.json" + jdata = j_loader(jfile) + + 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 + + jdata["model"]["descriptor"].pop("type", None) + jdata["model"]["descriptor"]["ntypes"] = 2 + jdata["model"]["descriptor"]["stripped_type_embedding"] = True + jdata["model"]["descriptor"]["attn_layer"] = 2 + descrpt = DescrptSeAtten(**jdata["model"]["descriptor"], uniform_seed=True) + jdata["model"]["fitting_net"]["descrpt"] = descrpt + fitting = EnerFitting(**jdata["model"]["fitting_net"], uniform_seed=True) + typeebd_param = jdata["model"]["type_embedding"] + typeebd = TypeEmbedNet( + neuron=typeebd_param["neuron"], + activation_function=None, + resnet_dt=typeebd_param["resnet_dt"], + seed=typeebd_param["seed"], + uniform_seed=True, + padding=True, + ) + model = EnerModel(descrpt, fitting, typeebd) + + # 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 + inputs_dict = {} + + model_pred = model.build( + t_coord, + t_type, + t_natoms, + t_box, + t_mesh, + inputs_dict, + suffix=self.filename + + "-" + + inspect.stack()[0][3] + + "test_model_se_atten_model_compressible", + 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 = self.test_session().__enter__() + sess.run(tf.global_variables_initializer()) + [e, f, v] = sess.run([energy, force, virial], feed_dict=feed_dict_test) + # print(sess.run(model.type_embedding)) + # np.savetxt('tmp.out', sess.run(descrpt.dout, feed_dict = feed_dict_test), fmt='%.10e') + # # print(sess.run(model.atype_embed, feed_dict = feed_dict_test)) + # print(sess.run(fitting.inputs, feed_dict = feed_dict_test)) + # print(sess.run(fitting.outs, feed_dict = feed_dict_test)) + # print(sess.run(fitting.atype_embed, feed_dict = feed_dict_test)) + + e = e.reshape([-1]) + f = f.reshape([-1]) + v = v.reshape([-1]) + np.savetxt("e.out", e.reshape([1, -1]), delimiter=",") + np.savetxt("f.out", f.reshape([1, -1]), delimiter=",") + np.savetxt("v.out", v.reshape([1, -1]), delimiter=",") + + refe = [6.125926357944699419e01] + reff = [ + 4.071033392194846855e-03, + 1.191078506555811808e-02, + 5.710038490045591959e-04, + 2.083813511902148086e-02, + 1.050404909007916256e-02, + -1.935131519230624082e-03, + -1.844253334357196135e-03, + 7.073208513688628192e-03, + -5.000418009101099666e-05, + 2.877594036828151017e-03, + -1.849276075329028823e-02, + -5.424378676202407318e-04, + -3.237425532982485255e-02, + 2.747768700765259881e-03, + 2.122946741188093227e-03, + 6.431746116137571252e-03, + -1.374305061680087571e-02, + -1.663770232507767613e-04, + ] + refv = [ + -5.699812217912397089e-02, + 8.904976757403390217e-03, + 2.306167440955633578e-03, + 8.904976757403417972e-03, + -3.502855693434053092e-02, + -6.596869271547715083e-04, + 2.306167440955633578e-03, + -6.596869271547720504e-04, + -1.602012510288682717e-04, + ] + + refe = np.reshape(refe, [-1]) + reff = np.reshape(reff, [-1]) + refv = np.reshape(refv, [-1]) + + places = 10 + np.testing.assert_almost_equal(e, refe, places) + np.testing.assert_almost_equal(f, reff, places) + np.testing.assert_almost_equal(v, refv, places) + + def test_stripped_type_embedding_exclude_types(self): + """In this test, we make type 0 has no interaction with type 0 and type 1, + so the descriptor should be zero for type 0 atoms. + """ + jfile = "water_se_atten.json" + jdata = j_loader(jfile) + + 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["model"]["descriptor"], "rcut") + ntypes = 2 + + data = DataSystem(systems, set_pfx, batch_size, test_size, rcut, run_opt=None) + + test_data = data.get_test() + numb_test = 1 + + # set parameters + jdata["model"]["descriptor"]["exclude_types"] = [[0, 0], [0, 1]] + + t_prop_c = tf.placeholder(tf.float32, [5], name="t_prop_c") + 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, [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) + + # successful + descrpt = DescrptSeAtten(ntypes=ntypes, **jdata["model"]["descriptor"]) + typeebd_param = jdata["model"]["type_embedding"] + jdata["model"]["descriptor"]["stripped_type_embedding"] = True + jdata["model"]["descriptor"]["attn_layer"] = 2 + typeebd = TypeEmbedNet( + neuron=typeebd_param["neuron"], + activation_function=None, + resnet_dt=typeebd_param["resnet_dt"], + seed=typeebd_param["seed"], + uniform_seed=True, + padding=True, + ) + type_embedding = typeebd.build( + ntypes, + suffix=self.filename + "-" + inspect.stack()[0][3] + "_type_embedding", + ) + dout = descrpt.build( + t_coord, + t_type, + t_natoms, + t_box, + t_mesh, + {"type_embedding": type_embedding}, + reuse=False, + suffix=self.filename + + "-" + + inspect.stack()[0][3] + + "_se_atten_compressible_exclude_types", + ) + + feed_dict_test1 = { + 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, + } + + with self.test_session() as sess: + sess.run(tf.global_variables_initializer()) + [des] = sess.run([dout], feed_dict=feed_dict_test1) + + np.testing.assert_almost_equal(des[:, 0:2], 0.0, 10) + with self.assertRaises(AssertionError): + np.testing.assert_almost_equal(des[:, 2:6], 0.0, 10) diff --git a/source/tests/water_se_atten_compressible_mixed_type.json b/source/tests/water_se_atten_compressible_mixed_type.json new file mode 100644 index 0000000000..04c84e3160 --- /dev/null +++ b/source/tests/water_se_atten_compressible_mixed_type.json @@ -0,0 +1,80 @@ +{ + "_comment": " model parameters", + "model": { + "type_map": [ + "foo", + "bar" + ], + "type_embedding": { + "neuron": [ + 8 + ], + "resnet_dt": false, + "seed": 1 + }, + "descriptor": { + "type": "se_atten", + "compressible": true, + "sel": 120, + "rcut_smth": 5.80, + "rcut": 6.00, + "neuron": [ + 25, + 50, + 100 + ], + "resnet_dt": false, + "type_one_side": false, + "axis_neuron": 16, + "seed": 1, + "attn": 128, + "attn_layer": 0, + "attn_dotr": true, + "attn_mask": false + }, + "fitting_net": { + "neuron": [ + 240, + 240, + 240 + ], + "resnet_dt": true, + "seed": 1 + } + }, + + "_comment": " traing controls", + "systems": [ + "system_mixed_type" + ], + "set_prefix": "set", + "stop_batch": 1000000, + "batch_size": 1, + "start_lr": 0.005, + "decay_steps": 5000, + "decay_rate": 0.95, + + "start_pref_e": 0.02, + "limit_pref_e": 1, + "start_pref_f": 1000, + "limit_pref_f": 1, + "start_pref_v": 0, + "limit_pref_v": 0, + + "seed": 1, + + "_comment": " display and restart", + "_comment": " frequencies counted in batch", + "disp_file": "lcurve.out", + "disp_freq": 100, + "numb_test": 1, + "save_freq": 1000, + "save_ckpt": "model.ckpt", + "load_ckpt": "model.ckpt", + "disp_training": true, + "time_training": true, + "profiling": false, + "profiling_file": "timeline.json", + + "_comment": "that's all" +}