From 71dd24acd0680166e601cbc5d75d577937d6403a Mon Sep 17 00:00:00 2001 From: nahso Date: Thu, 14 Sep 2023 17:57:43 +0800 Subject: [PATCH 01/10] apply compression for se_e2_a_tebd --- deepmd/descriptor/__init__.py | 4 + deepmd/descriptor/se_a.py | 283 +++++++++++++++++++++++++++++++--- deepmd/descriptor/se_a_v2.py | 67 ++++++++ deepmd/descriptor/se_atten.py | 1 + deepmd/utils/argcheck.py | 5 + deepmd/utils/tabulate.py | 10 +- 6 files changed, 341 insertions(+), 29 deletions(-) create mode 100644 deepmd/descriptor/se_a_v2.py diff --git a/deepmd/descriptor/__init__.py b/deepmd/descriptor/__init__.py index ab726d95c8..6bb6d4a09d 100644 --- a/deepmd/descriptor/__init__.py +++ b/deepmd/descriptor/__init__.py @@ -11,6 +11,9 @@ from .se_a import ( DescrptSeA, ) +from .se_a_v2 import ( + DescrptSeAV2, +) from .se_a_ebd import ( DescrptSeAEbd, ) @@ -39,6 +42,7 @@ "DescrptHybrid", "DescrptLocFrame", "DescrptSeA", + "DescrptSeAV2", "DescrptSeAEbd", "DescrptSeAEf", "DescrptSeAEfLower", diff --git a/deepmd/descriptor/se_a.py b/deepmd/descriptor/se_a.py index 1349f61464..7b985deb38 100644 --- a/deepmd/descriptor/se_a.py +++ b/deepmd/descriptor/se_a.py @@ -11,6 +11,7 @@ cast_precision, get_activation_func, get_precision, + get_np_precision, ) from deepmd.env import ( GLOBAL_NP_FLOAT_PRECISION, @@ -52,6 +53,12 @@ from deepmd.utils.type_embed import ( embed_atom_type, ) +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 .descriptor import ( Descriptor, @@ -165,6 +172,7 @@ def __init__( uniform_seed: bool = False, multi_task: bool = False, spin: Optional[Spin] = None, + stripped_type_embedding: bool = False, **kwargs, ) -> None: """Constructor.""" @@ -185,6 +193,7 @@ def __init__( self.compress_activation_fn = get_activation_func(activation_function) self.filter_activation_fn = get_activation_func(activation_function) self.filter_precision = get_precision(precision) + self.filter_np_precision = get_np_precision(precision) self.exclude_types = set() for tt in exclude_types: assert len(tt) == 2 @@ -193,6 +202,9 @@ def __init__( self.set_davg_zero = set_davg_zero self.type_one_side = type_one_side self.spin = spin + self.stripped_type_embedding = stripped_type_embedding + self.extra_embeeding_net_variables = None + self.layer_size = len(neuron) # extend sel_a for spin system if self.spin is not None: @@ -463,6 +475,35 @@ def enable_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.stripped_type_embedding: + ret_two_side = get_pattern_nodes_from_graph_def( + graph_def, f"filter_type_all{suffix}/.+_two_side_ebd" + ) + ret_one_side = get_pattern_nodes_from_graph_def( + graph_def, f"filter_type_all{suffix}/.+_one_side_ebd" + ) + if len(ret_two_side) == 0 and len(ret_one_side) == 0: + raise RuntimeError( + "can not find variables of embedding net from graph_def, maybe it is not a compressible model." + ) + elif len(ret_one_side) != 0 and len(ret_two_side) != 0: + raise RuntimeError( + "both one side and two side embedding net varaibles are detected, it is a wrong model." + ) + elif len(ret_two_side) != 0: + self.final_type_embedding = self._get_two_side_type_embedding(graph) + self.matrix = self._get_extra_side_embedding_net_variable( + graph_def, "two_side", "matrix", suffix + ) + self.bias = self._get_extra_side_embedding_net_variable(graph_def, "two_side", "bias", suffix) + self.extra_embedding = self._make_data(self.final_type_embedding) + else: + self.final_type_embedding = self._get_type_embedding(graph) + self.matrix = self._get_extra_side_embedding_net_variable( + graph_def, "one_side", "matrix", suffix + ) + self.bias = self._get_extra_side_embedding_net_variable(graph_def, "one_side", "bias", suffix) + self.extra_embedding = self._make_data(self.final_type_embedding) self.compress = True self.table = DPTabulate( @@ -588,6 +629,7 @@ def build( coord = tf.reshape(coord_, [-1, natoms[1] * 3]) box = tf.reshape(box_, [-1, 9]) atype = tf.reshape(atype_, [-1, natoms[1]]) + self.atype = atype op_descriptor = ( build_op_descriptor() if nvnmd_cfg.enable else op_module.prod_env_mat_a @@ -606,6 +648,10 @@ def build( sel_a=self.sel_a, sel_r=self.sel_r, ) + nlist_t = tf.reshape(self.nlist + 1, [-1]) + atype_t = tf.concat([[self.ntypes], tf.reshape(self.atype, [-1])], axis=0) + self.nei_type_vec = tf.nn.embedding_lookup(atype_t, nlist_t) + # only used when tensorboard was set as true tf.summary.histogram("descrpt", self.descrpt) tf.summary.histogram("rij", self.rij) @@ -901,13 +947,89 @@ 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]) if type_embedding is not None: - xyz_scatter = self._concat_type_embedding( - xyz_scatter, nframes, natoms, type_embedding - ) - if self.compress: - raise RuntimeError( - "compression of type embedded descriptor is not supported at the moment" + if self.stripped_type_embedding: + if self.type_one_side: + extra_embedding_index = self.nei_type_vec + else: + padding_ntypes = type_embedding.shape[0] + atype_expand = tf.reshape(self.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]) + extra_embedding_index = index_of_two_side + + if not self.compress: + if self.type_one_side: + one_side_type_embedding_suffix = "_one_side_ebd" + net_output = embedding_net( + type_embedding, + self.filter_neuron, + self.filter_precision, + activation_fn=activation_fn, + resnet_dt=self.filter_resnet_dt, + name_suffix=one_side_type_embedding_suffix, + stddev=stddev, + bavg=bavg, + seed=self.seed, + trainable=trainable, + uniform_seed=self.uniform_seed, + initial_variables=self.extra_embeeding_net_variables, + mixed_prec=self.mixed_prec, + ) + net_output = tf.nn.embedding_lookup(net_output, self.nei_type_vec) + 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]], + ) + + atype_expand = tf.reshape(self.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]) + self.extra_embedding_index = index_of_two_side + + two_side_type_embedding_suffix = "_two_side_ebd" + net_output = 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.extra_embeeding_net_variables, + mixed_prec=self.mixed_prec, + ) + net_output = tf.nn.embedding_lookup( + net_output, idx + ) + net_output = tf.reshape(net_output, [-1, self.filter_neuron[-1]]) + else: + xyz_scatter = self._concat_type_embedding( + xyz_scatter, nframes, natoms, type_embedding ) + if self.compress: + raise RuntimeError( + "compression of type embedded descriptor is not supported when stripped_type_embedding == False" + ) # natom x 4 x outputs_size if nvnmd_cfg.enable: return filter_lower_R42GR( @@ -929,25 +1051,48 @@ def _filter_lower( self.embedding_net_variables, ) if self.compress and (not is_exclude): - if self.type_one_side: - net = "filter_-1_net_" + str(type_i) + if self.stripped_type_embedding: + net_output = tf.nn.embedding_lookup( + self.extra_embedding, extra_embedding_index + ) + 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], + ] + 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]), + net_output, + last_layer_size=outputs_size[-1], + is_sorted=False + ) else: - net = "filter_" + str(type_input) + "_net_" + str(type_i) - 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], - ] - return op_module.tabulate_fusion_se_a( - tf.cast(self.table.data[net], self.filter_precision), - info, - xyz_scatter, - tf.reshape(inputs_i, [natom, shape_i[1] // 4, 4]), - last_layer_size=outputs_size[-1], - ) + if self.type_one_side: + net = "filter_-1_net_" + str(type_i) + else: + net = "filter_" + str(type_input) + "_net_" + str(type_i) + 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], + ] + return op_module.tabulate_fusion_se_a( + tf.cast(self.table.data[net], self.filter_precision), + info, + xyz_scatter, + tf.reshape(inputs_i, [natom, shape_i[1] // 4, 4]), + last_layer_size=outputs_size[-1], + ) else: if not is_exclude: # with (natom x nei_type_i) x out_size @@ -966,6 +1111,9 @@ def _filter_lower( initial_variables=self.embedding_net_variables, mixed_prec=self.mixed_prec, ) + + if self.stripped_type_embedding: + xyz_scatter = xyz_scatter * net_output + xyz_scatter if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift else: @@ -1179,3 +1327,90 @@ def init_variables( self.dstd = new_dstd if self.original_sel is None: self.original_sel = sel + @property + def explicit_ntypes(self) -> bool: + """Explicit ntypes with type embedding.""" + print('inner se_a, explicit_ntypes = ', self.stripped_type_embedding) + if self.stripped_type_embedding: + return True + return False + + def _get_type_embedding(self, graph): + type_embedding = get_tensor_by_name_from_graph(graph, "t_typeebd") + type_embedding = type_embedding.astype(self.filter_np_precision) + return type_embedding + + def _get_two_side_type_embedding(self, graph): + # todo: refactor + 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_extra_side_embedding_net_variable(self, graph_def, type_side, 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}_{type_side}_ebd", + ) + node = target[f"filter_type_all{suffix}/{varialbe_name}_{i}_{type_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 diff --git a/deepmd/descriptor/se_a_v2.py b/deepmd/descriptor/se_a_v2.py new file mode 100644 index 0000000000..521a63ad1a --- /dev/null +++ b/deepmd/descriptor/se_a_v2.py @@ -0,0 +1,67 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import logging +from typing import ( + List, + Optional, +) + +from .descriptor import ( + Descriptor, +) +from .se_a import ( + DescrptSeA, +) +from deepmd.utils.spin import ( + Spin, +) + +log = logging.getLogger(__name__) + + +@Descriptor.register("se_e2_a_v2") +class DescrptSeAV2(DescrptSeA): + r""" + A warpper for DescriptorSeA, which set stripped_type_embedding=True + """ + + def __init__( + self, + rcut: float, + rcut_smth: float, + sel: List[str], + neuron: List[int] = [24, 48, 96], + axis_neuron: int = 8, + resnet_dt: bool = False, + trainable: bool = True, + seed: Optional[int] = None, + type_one_side: bool = True, + exclude_types: List[List[int]] = [], + set_davg_zero: bool = False, + activation_function: str = "tanh", + precision: str = "default", + uniform_seed: bool = False, + multi_task: bool = False, + spin: Optional[Spin] = None, + **kwargs, + ) -> None: + DescrptSeA.__init__( + self, + rcut, + rcut_smth, + sel, + neuron=neuron, + axis_neuron=axis_neuron, + resnet_dt=resnet_dt, + trainable=trainable, + seed=seed, + type_one_side=type_one_side, + exclude_types=exclude_types, + set_davg_zero=set_davg_zero, + activation_function=activation_function, + precision=precision, + uniform_seed=uniform_seed, + multi_task=multi_task, + spin=spin, + stripped_type_embedding=True, + **kwargs, + ) diff --git a/deepmd/descriptor/se_atten.py b/deepmd/descriptor/se_atten.py index 12558c45c4..ee51e1d141 100644 --- a/deepmd/descriptor/se_atten.py +++ b/deepmd/descriptor/se_atten.py @@ -429,6 +429,7 @@ def enable_compression( graph, "descrpt_attr%s/t_std" % suffix ) + # todo: refactor 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) diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 153824cb0d..46f2c7da0e 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -429,6 +429,11 @@ def descrpt_se_atten_v2_args(): ] +@descrpt_args_plugin.register("se_e2_a_v2") +def descrpt_se_e2_a_v2_args(): + return descrpt_se_a_args() + + @descrpt_args_plugin.register("se_a_mask") def descrpt_se_a_mask_args(): doc_sel = 'This parameter sets the number of selected neighbors for each type of atom. It can be:\n\n\ diff --git a/deepmd/utils/tabulate.py b/deepmd/utils/tabulate.py index 883730b9d9..81eacaa459 100644 --- a/deepmd/utils/tabulate.py +++ b/deepmd/utils/tabulate.py @@ -176,7 +176,7 @@ 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.DescrptSeAtten): + if isinstance(self.descrpt, deepmd.descriptor.DescrptSeAtten) or isinstance(self.descrpt, deepmd.descriptor.DescrptSeAV2): uu = np.max(upper) ll = np.min(lower) xx = np.arange(ll, uu, stride0, dtype=self.data_type) @@ -419,7 +419,7 @@ def _get_bias(self): bias = {} for layer in range(1, self.layer_size + 1): bias["layer_" + str(layer)] = [] - if isinstance(self.descrpt, deepmd.descriptor.DescrptSeAtten): + if isinstance(self.descrpt, deepmd.descriptor.DescrptSeAtten) or isinstance(self.descrpt, deepmd.descriptor.DescrptSeAV2): node = self.embedding_net_nodes[ f"filter_type_all{self.suffix}/bias_{layer}" ] @@ -483,7 +483,7 @@ def _get_matrix(self): matrix = {} for layer in range(1, self.layer_size + 1): matrix["layer_" + str(layer)] = [] - if isinstance(self.descrpt, deepmd.descriptor.DescrptSeAtten): + if isinstance(self.descrpt, deepmd.descriptor.DescrptSeAtten) or isinstance(self.descrpt, deepmd.descriptor.DescrptSeAV2): node = self.embedding_net_nodes[ f"filter_type_all{self.suffix}/matrix_{layer}" ] @@ -687,7 +687,7 @@ def _spline5_switch(self, xx, rmin, rmax): def _get_layer_size(self): layer_size = 0 - if isinstance(self.descrpt, deepmd.descriptor.DescrptSeAtten): + if isinstance(self.descrpt, deepmd.descriptor.DescrptSeAtten) or isinstance(self.descrpt, deepmd.descriptor.DescrptSeAV2): layer_size = len(self.embedding_net_nodes) // 2 elif isinstance(self.descrpt, deepmd.descriptor.DescrptSeA): layer_size = len(self.embedding_net_nodes) // ( @@ -737,7 +737,7 @@ def _all_excluded(self, ii: int) -> bool: def _get_table_size(self): table_size = 0 - if isinstance(self.descrpt, deepmd.descriptor.DescrptSeAtten): + if isinstance(self.descrpt, deepmd.descriptor.DescrptSeAtten) or isinstance(self.descrpt, deepmd.descriptor.DescrptSeAV2): table_size = 1 elif isinstance(self.descrpt, deepmd.descriptor.DescrptSeA): table_size = self.ntypes * self.ntypes From 8e36876199473b4b6fac7d7d850807dfafb2a26e Mon Sep 17 00:00:00 2001 From: nahso Date: Tue, 19 Sep 2023 13:55:00 +0800 Subject: [PATCH 02/10] ut for se_e2_a_ebd_v2 --- deepmd/descriptor/se_a_mask.py | 2 + .../tests/test_model_compression_se_a_ebd.py | 587 ++++++++++++++++++ ...odel_compression_se_a_ebd_type_one_side.py | 585 +++++++++++++++++ 3 files changed, 1174 insertions(+) create mode 100644 source/tests/test_model_compression_se_a_ebd.py create mode 100644 source/tests/test_model_compression_se_a_ebd_type_one_side.py diff --git a/deepmd/descriptor/se_a_mask.py b/deepmd/descriptor/se_a_mask.py index cdec33e292..2eaf09d035 100644 --- a/deepmd/descriptor/se_a_mask.py +++ b/deepmd/descriptor/se_a_mask.py @@ -128,6 +128,7 @@ def __init__( activation_function: str = "tanh", precision: str = "default", uniform_seed: bool = False, + stripped_type_embedding: bool = False, **kwargs, ) -> None: """Constructor.""" @@ -159,6 +160,7 @@ def __init__( # numb of neighbors and numb of descrptors self.nnei_a = np.cumsum(self.sel_a)[-1] self.nnei = self.nnei_a + self.stripped_type_embedding = stripped_type_embedding self.ndescrpt_a = self.nnei_a * 4 self.ndescrpt = self.ndescrpt_a diff --git a/source/tests/test_model_compression_se_a_ebd.py b/source/tests/test_model_compression_se_a_ebd.py new file mode 100644 index 0000000000..78ef639a11 --- /dev/null +++ b/source/tests/test_model_compression_se_a_ebd.py @@ -0,0 +1,587 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +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 deepmd.env import ( + GLOBAL_NP_FLOAT_PRECISION, +) +from deepmd.infer import ( + DeepPot, +) + +if GLOBAL_NP_FLOAT_PRECISION == np.float32: + default_places = 4 +else: + default_places = 10 + + +def _file_delete(file): + if os.path.isdir(file): + os.rmdir(file) + elif os.path.isfile(file): + os.remove(file) + + +def _subprocess_run(command): + popen = sp.Popen(command.split(), shell=False, stdout=sp.PIPE, stderr=sp.STDOUT) + for line in iter(popen.stdout.readline, b""): + if hasattr(line, "decode"): + line = line.decode("utf-8") + line = line.rstrip() + print(line) + popen.wait() + return popen.returncode + + +def _init_models(): + data_file = str(tests_path / os.path.join("model_compression", "data")) + frozen_model = str(tests_path / "dp-original-se-e2-a-v2.pb") + compressed_model = str(tests_path / "dp-compressed-se-e2-a-v2.pb") + INPUT = str(tests_path / "input_se_e2_a_v2.json") + jdata = j_loader(str(tests_path / os.path.join("model_compression", "input.json"))) + jdata["training"]["training_data"]["systems"] = data_file + jdata["training"]["validation_data"]["systems"] = data_file + jdata["model"]["descriptor"] = {} + jdata["model"]["descriptor"]["type"] = "se_e2_a_v2" + jdata["model"]["descriptor"]["type_one_side"] = False + jdata["model"]["type_embedding"] = {} + jdata["model"]["type_embedding"]["neuron"] = [1] + with open(INPUT, "w") as fp: + json.dump(jdata, fp, indent=4) + with open("input_v2_compat.json.tempfile", "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!") + return INPUT, frozen_model, compressed_model + + +def _init_models_exclude_types(): + data_file = str(tests_path / os.path.join("model_compression", "data")) + frozen_model = str(tests_path / "dp-original-exclude-types-se-e2-a-v2.pb") + compressed_model = str(tests_path / "dp-compressed-exclude-types-se-e2-a-v2.pb") + INPUT = str(tests_path / "input_se_e2_a_v2.json") + jdata = j_loader(str(tests_path / os.path.join("model_compression", "input.json"))) + jdata["model"]["descriptor"] = {} + jdata["model"]["descriptor"]["type"] = "se_e2_a_v2" + jdata["model"]["descriptor"]["exclude_types"] = [[0, 1]] + jdata["model"]["descriptor"]["type_one_side"] = False + jdata["model"]["type_embedding"] = {} + jdata["model"]["type_embedding"]["neuron"] = [1] + 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!") + return INPUT, frozen_model, compressed_model + + +INPUT, FROZEN_MODEL, COMPRESSED_MODEL = _init_models() +INPUT_ET, FROZEN_MODEL_ET, COMPRESSED_MODEL_ET = _init_models_exclude_types() + + +class TestDeepPotAPBC(unittest.TestCase): + @classmethod + def setUpClass(self): + self.dp_original = DeepPot(FROZEN_MODEL) + self.dp_compressed = DeepPot(COMPRESSED_MODEL) + 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): + self.assertEqual(self.dp_original.get_ntypes(), 2) + self.assertAlmostEqual(self.dp_original.get_rcut(), 6.0, places=default_places) + self.assertEqual(self.dp_original.get_type_map(), ["O", "H"]) + self.assertEqual(self.dp_original.get_dim_fparam(), 0) + self.assertEqual(self.dp_original.get_dim_aparam(), 0) + + self.assertEqual(self.dp_compressed.get_ntypes(), 2) + self.assertAlmostEqual( + self.dp_compressed.get_rcut(), 6.0, places=default_places + ) + self.assertEqual(self.dp_compressed.get_type_map(), ["O", "H"]) + self.assertEqual(self.dp_compressed.get_dim_fparam(), 0) + self.assertEqual(self.dp_compressed.get_dim_aparam(), 0) + + def test_1frame(self): + ee0, ff0, vv0 = self.dp_original.eval( + self.coords, self.box, self.atype, atomic=False + ) + ee1, ff1, vv1 = self.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): + ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( + self.coords, self.box, self.atype, atomic=True + ) + ee1, ff1, vv1, ae1, av1 = self.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): + coords2 = np.concatenate((self.coords, self.coords)) + box2 = np.concatenate((self.box, self.box)) + ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( + coords2, box2, self.atype, atomic=True + ) + ee1, ff1, vv1, ae1, av1 = self.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) + + +class TestDeepPotANoPBC(unittest.TestCase): + @classmethod + def setUpClass(self): + self.dp_original = DeepPot(FROZEN_MODEL) + self.dp_compressed = DeepPot(COMPRESSED_MODEL) + 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): + ee0, ff0, vv0 = self.dp_original.eval( + self.coords, self.box, self.atype, atomic=False + ) + ee1, ff1, vv1 = self.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): + ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( + self.coords, self.box, self.atype, atomic=True + ) + ee1, ff1, vv1, ae1, av1 = self.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): + coords2 = np.concatenate((self.coords, self.coords)) + ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( + coords2, self.box, self.atype, atomic=True + ) + ee1, ff1, vv1, ae1, av1 = self.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) + + +class TestDeepPotALargeBoxNoPBC(unittest.TestCase): + @classmethod + def setUpClass(self): + self.dp_original = DeepPot(FROZEN_MODEL) + self.dp_compressed = DeepPot(COMPRESSED_MODEL) + 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): + ee0, ff0, vv0 = self.dp_original.eval( + self.coords, self.box, self.atype, atomic=False + ) + ee1, ff1, vv1 = self.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): + ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( + self.coords, self.box, self.atype, atomic=True + ) + ee1, ff1, vv1, ae1, av1 = self.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): + 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_MODEL), + ) + water1 = Atoms( + "OHHOHH", + positions=self.coords.reshape((-1, 3)), + cell=self.box.reshape((3, 3)), + calculator=DP(COMPRESSED_MODEL), + ) + 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) + + +class TestDeepPotAPBCExcludeTypes(unittest.TestCase): + @classmethod + def setUpClass(self): + self.dp_original = DeepPot(FROZEN_MODEL_ET) + self.dp_compressed = DeepPot(COMPRESSED_MODEL_ET) + 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): + _file_delete(INPUT_ET) + _file_delete(FROZEN_MODEL_ET) + _file_delete(COMPRESSED_MODEL_ET) + _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): + self.assertEqual(self.dp_original.get_ntypes(), 2) + self.assertAlmostEqual(self.dp_original.get_rcut(), 6.0, places=default_places) + self.assertEqual(self.dp_original.get_type_map(), ["O", "H"]) + self.assertEqual(self.dp_original.get_dim_fparam(), 0) + self.assertEqual(self.dp_original.get_dim_aparam(), 0) + + self.assertEqual(self.dp_compressed.get_ntypes(), 2) + self.assertAlmostEqual( + self.dp_compressed.get_rcut(), 6.0, places=default_places + ) + self.assertEqual(self.dp_compressed.get_type_map(), ["O", "H"]) + self.assertEqual(self.dp_compressed.get_dim_fparam(), 0) + self.assertEqual(self.dp_compressed.get_dim_aparam(), 0) + + def test_1frame(self): + ee0, ff0, vv0 = self.dp_original.eval( + self.coords, self.box, self.atype, atomic=False + ) + ee1, ff1, vv1 = self.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): + ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( + self.coords, self.box, self.atype, atomic=True + ) + ee1, ff1, vv1, ae1, av1 = self.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): + coords2 = np.concatenate((self.coords, self.coords)) + box2 = np.concatenate((self.box, self.box)) + ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( + coords2, box2, self.atype, atomic=True + ) + ee1, ff1, vv1, ae1, av1 = self.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_compression_se_a_ebd_type_one_side.py b/source/tests/test_model_compression_se_a_ebd_type_one_side.py new file mode 100644 index 0000000000..9787a4e68a --- /dev/null +++ b/source/tests/test_model_compression_se_a_ebd_type_one_side.py @@ -0,0 +1,585 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +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 deepmd.env import ( + GLOBAL_NP_FLOAT_PRECISION, +) +from deepmd.infer import ( + DeepPot, +) + +if GLOBAL_NP_FLOAT_PRECISION == np.float32: + default_places = 4 +else: + default_places = 10 + + +def _file_delete(file): + if os.path.isdir(file): + os.rmdir(file) + elif os.path.isfile(file): + os.remove(file) + + +def _subprocess_run(command): + popen = sp.Popen(command.split(), shell=False, stdout=sp.PIPE, stderr=sp.STDOUT) + for line in iter(popen.stdout.readline, b""): + if hasattr(line, "decode"): + line = line.decode("utf-8") + line = line.rstrip() + print(line) + popen.wait() + return popen.returncode + + +def _init_models(): + data_file = str(tests_path / os.path.join("model_compression", "data")) + frozen_model = str(tests_path / "dp-original-se-e2-a-v2-one-side.pb") + compressed_model = str(tests_path / "dp-compressed-se-e2-a-v2-one-side.pb") + INPUT = str(tests_path / "input.json") + jdata = j_loader(str(tests_path / os.path.join("model_compression", "input.json"))) + jdata["training"]["training_data"]["systems"] = data_file + jdata["training"]["validation_data"]["systems"] = data_file + jdata["model"]["descriptor"] = {} + jdata["model"]["descriptor"]["type"] = "se_e2_a_v2" + jdata["model"]["descriptor"]["type_one_side"] = True + jdata["model"]["type_embedding"] = {} + jdata["model"]["type_embedding"]["neuron"] = [1] + 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!") + return INPUT, frozen_model, compressed_model + + +def _init_models_exclude_types(): + data_file = str(tests_path / os.path.join("model_compression", "data")) + frozen_model = str(tests_path / "dp-original-exclude-types-se-e2-a-v2-one-side.pb") + compressed_model = str(tests_path / "dp-compressed-exclude-types-se-e2-a-v2-one-side.pb") + INPUT = str(tests_path / "input.json") + jdata = j_loader(str(tests_path / os.path.join("model_compression", "input.json"))) + jdata["model"]["descriptor"] = {} + jdata["model"]["descriptor"]["type"] = "se_e2_a_v2" + jdata["model"]["descriptor"]["exclude_types"] = [[0, 1]] + jdata["model"]["descriptor"]["type_one_side"] = True + jdata["model"]["type_embedding"] = {} + jdata["model"]["type_embedding"]["neuron"] = [1] + 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!") + return INPUT, frozen_model, compressed_model + + +INPUT, FROZEN_MODEL, COMPRESSED_MODEL = _init_models() +INPUT_ET, FROZEN_MODEL_ET, COMPRESSED_MODEL_ET = _init_models_exclude_types() + + +class TestDeepPotAPBC(unittest.TestCase): + @classmethod + def setUpClass(self): + self.dp_original = DeepPot(FROZEN_MODEL) + self.dp_compressed = DeepPot(COMPRESSED_MODEL) + 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): + self.assertEqual(self.dp_original.get_ntypes(), 2) + self.assertAlmostEqual(self.dp_original.get_rcut(), 6.0, places=default_places) + self.assertEqual(self.dp_original.get_type_map(), ["O", "H"]) + self.assertEqual(self.dp_original.get_dim_fparam(), 0) + self.assertEqual(self.dp_original.get_dim_aparam(), 0) + + self.assertEqual(self.dp_compressed.get_ntypes(), 2) + self.assertAlmostEqual( + self.dp_compressed.get_rcut(), 6.0, places=default_places + ) + self.assertEqual(self.dp_compressed.get_type_map(), ["O", "H"]) + self.assertEqual(self.dp_compressed.get_dim_fparam(), 0) + self.assertEqual(self.dp_compressed.get_dim_aparam(), 0) + + def test_1frame(self): + ee0, ff0, vv0 = self.dp_original.eval( + self.coords, self.box, self.atype, atomic=False + ) + ee1, ff1, vv1 = self.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): + ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( + self.coords, self.box, self.atype, atomic=True + ) + ee1, ff1, vv1, ae1, av1 = self.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): + coords2 = np.concatenate((self.coords, self.coords)) + box2 = np.concatenate((self.box, self.box)) + ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( + coords2, box2, self.atype, atomic=True + ) + ee1, ff1, vv1, ae1, av1 = self.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) + + +class TestDeepPotANoPBC(unittest.TestCase): + @classmethod + def setUpClass(self): + self.dp_original = DeepPot(FROZEN_MODEL) + self.dp_compressed = DeepPot(COMPRESSED_MODEL) + 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): + ee0, ff0, vv0 = self.dp_original.eval( + self.coords, self.box, self.atype, atomic=False + ) + ee1, ff1, vv1 = self.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): + ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( + self.coords, self.box, self.atype, atomic=True + ) + ee1, ff1, vv1, ae1, av1 = self.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): + coords2 = np.concatenate((self.coords, self.coords)) + ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( + coords2, self.box, self.atype, atomic=True + ) + ee1, ff1, vv1, ae1, av1 = self.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) + + +class TestDeepPotALargeBoxNoPBC(unittest.TestCase): + @classmethod + def setUpClass(self): + self.dp_original = DeepPot(FROZEN_MODEL) + self.dp_compressed = DeepPot(COMPRESSED_MODEL) + 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): + ee0, ff0, vv0 = self.dp_original.eval( + self.coords, self.box, self.atype, atomic=False + ) + ee1, ff1, vv1 = self.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): + ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( + self.coords, self.box, self.atype, atomic=True + ) + ee1, ff1, vv1, ae1, av1 = self.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): + 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_MODEL), + ) + water1 = Atoms( + "OHHOHH", + positions=self.coords.reshape((-1, 3)), + cell=self.box.reshape((3, 3)), + calculator=DP(COMPRESSED_MODEL), + ) + 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) + + +class TestDeepPotAPBCExcludeTypes(unittest.TestCase): + @classmethod + def setUpClass(self): + self.dp_original = DeepPot(FROZEN_MODEL_ET) + self.dp_compressed = DeepPot(COMPRESSED_MODEL_ET) + 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): + _file_delete(INPUT_ET) + _file_delete(FROZEN_MODEL_ET) + _file_delete(COMPRESSED_MODEL_ET) + _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): + self.assertEqual(self.dp_original.get_ntypes(), 2) + self.assertAlmostEqual(self.dp_original.get_rcut(), 6.0, places=default_places) + self.assertEqual(self.dp_original.get_type_map(), ["O", "H"]) + self.assertEqual(self.dp_original.get_dim_fparam(), 0) + self.assertEqual(self.dp_original.get_dim_aparam(), 0) + + self.assertEqual(self.dp_compressed.get_ntypes(), 2) + self.assertAlmostEqual( + self.dp_compressed.get_rcut(), 6.0, places=default_places + ) + self.assertEqual(self.dp_compressed.get_type_map(), ["O", "H"]) + self.assertEqual(self.dp_compressed.get_dim_fparam(), 0) + self.assertEqual(self.dp_compressed.get_dim_aparam(), 0) + + def test_1frame(self): + ee0, ff0, vv0 = self.dp_original.eval( + self.coords, self.box, self.atype, atomic=False + ) + ee1, ff1, vv1 = self.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): + ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( + self.coords, self.box, self.atype, atomic=True + ) + ee1, ff1, vv1, ae1, av1 = self.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): + coords2 = np.concatenate((self.coords, self.coords)) + box2 = np.concatenate((self.box, self.box)) + ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( + coords2, box2, self.atype, atomic=True + ) + ee1, ff1, vv1, ae1, av1 = self.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) From c4ddb98c26548f1e5211971e6302aca9598257b7 Mon Sep 17 00:00:00 2001 From: nahso Date: Tue, 19 Sep 2023 15:05:11 +0800 Subject: [PATCH 03/10] refactor --- deepmd/descriptor/se_a.py | 105 ++++++---------------------------- deepmd/descriptor/se_atten.py | 90 ++++------------------------- deepmd/utils/compress.py | 96 +++++++++++++++++++++++++++++++ 3 files changed, 123 insertions(+), 168 deletions(-) create mode 100644 deepmd/utils/compress.py diff --git a/deepmd/descriptor/se_a.py b/deepmd/descriptor/se_a.py index 7b985deb38..66587b1466 100644 --- a/deepmd/descriptor/se_a.py +++ b/deepmd/descriptor/se_a.py @@ -59,6 +59,12 @@ get_tensor_by_name_from_graph, get_tensor_by_type, ) +from deepmd.utils.compress import ( + get_type_embedding, + get_two_side_type_embedding, + get_extra_side_embedding_net_variable, + make_data, +) from .descriptor import ( Descriptor, @@ -491,19 +497,19 @@ def enable_compression( "both one side and two side embedding net varaibles are detected, it is a wrong model." ) elif len(ret_two_side) != 0: - self.final_type_embedding = self._get_two_side_type_embedding(graph) - self.matrix = self._get_extra_side_embedding_net_variable( - graph_def, "two_side", "matrix", suffix + self.final_type_embedding = get_two_side_type_embedding(self, graph) + self.matrix = get_extra_side_embedding_net_variable( + self, graph_def, "two_side", "matrix", suffix ) - self.bias = self._get_extra_side_embedding_net_variable(graph_def, "two_side", "bias", suffix) - self.extra_embedding = self._make_data(self.final_type_embedding) + self.bias = get_extra_side_embedding_net_variable(self, graph_def, "two_side", "bias", suffix) + self.extra_embedding = make_data(self, self.final_type_embedding) else: - self.final_type_embedding = self._get_type_embedding(graph) - self.matrix = self._get_extra_side_embedding_net_variable( - graph_def, "one_side", "matrix", suffix + self.final_type_embedding = get_type_embedding(self, graph) + self.matrix = get_extra_side_embedding_net_variable( + self, graph_def, "one_side", "matrix", suffix ) - self.bias = self._get_extra_side_embedding_net_variable(graph_def, "one_side", "bias", suffix) - self.extra_embedding = self._make_data(self.final_type_embedding) + self.bias = get_extra_side_embedding_net_variable(self, graph_def, "one_side", "bias", suffix) + self.extra_embedding = make_data(self, self.final_type_embedding) self.compress = True self.table = DPTabulate( @@ -1335,82 +1341,3 @@ def explicit_ntypes(self) -> bool: return True return False - def _get_type_embedding(self, graph): - type_embedding = get_tensor_by_name_from_graph(graph, "t_typeebd") - type_embedding = type_embedding.astype(self.filter_np_precision) - return type_embedding - - def _get_two_side_type_embedding(self, graph): - # todo: refactor - 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_extra_side_embedding_net_variable(self, graph_def, type_side, 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}_{type_side}_ebd", - ) - node = target[f"filter_type_all{suffix}/{varialbe_name}_{i}_{type_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 diff --git a/deepmd/descriptor/se_atten.py b/deepmd/descriptor/se_atten.py index ee51e1d141..bc99847219 100644 --- a/deepmd/descriptor/se_atten.py +++ b/deepmd/descriptor/se_atten.py @@ -51,6 +51,12 @@ from deepmd.utils.tabulate import ( DPTabulate, ) +from deepmd.utils.compress import ( + get_type_embedding, + get_two_side_type_embedding, + get_extra_side_embedding_net_variable, + make_data, +) from .descriptor import ( Descriptor, @@ -415,12 +421,12 @@ def enable_compression( 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.final_type_embedding = get_two_side_type_embedding(self, graph) + self.matrix = get_extra_side_embedding_net_variable( + self, graph_def, "two_side", "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.bias = get_extra_side_embedding_net_variable(self, graph_def, "two_side", "bias", suffix) + self.two_embd = make_data(self, self.final_type_embedding) self.davg = get_tensor_by_name_from_graph( graph, "descrpt_attr%s/t_avg" % suffix @@ -429,80 +435,6 @@ def enable_compression( graph, "descrpt_attr%s/t_std" % suffix ) - # todo: refactor - 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}_two_side_ebd", - ) - node = target[f"filter_type_all{suffix}/{varialbe_name}_{i}_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, diff --git a/deepmd/utils/compress.py b/deepmd/utils/compress.py new file mode 100644 index 0000000000..7e394ed897 --- /dev/null +++ b/deepmd/utils/compress.py @@ -0,0 +1,96 @@ +import logging + +from deepmd.env import ( + tf, +) + +from deepmd.utils.graph import ( + get_pattern_nodes_from_graph_def, + get_tensor_by_name_from_graph, +) +import numpy as np + +log = logging.getLogger(__name__) + +def get_type_embedding(self, graph): + type_embedding = get_tensor_by_name_from_graph(graph, "t_typeebd") + type_embedding = type_embedding.astype(self.filter_np_precision) + return type_embedding + +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_extra_side_embedding_net_variable(self, graph_def, type_side, 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}_{type_side}_ebd", + ) + node = target[f"filter_type_all{suffix}/{varialbe_name}_{i}_{type_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 = ( + _layer_0( + self, + xx, + self.matrix["layer_" + str(layer + 1)], + self.bias["layer_" + str(layer + 1)], + ) + + xx + ) + elif self.filter_neuron[0] == 2: + tt, yy = _layer_1( + self, + xx, + self.matrix["layer_" + str(layer + 1)], + self.bias["layer_" + str(layer + 1)], + ) + else: + yy = _layer_0( + self, + xx, + self.matrix["layer_" + str(layer + 1)], + self.bias["layer_" + str(layer + 1)], + ) + else: + tt, zz = _layer_1( + self, + yy, + self.matrix["layer_" + str(layer + 1)], + self.bias["layer_" + str(layer + 1)], + ) + yy = zz + vv = sess.run(zz) + return vv From 937fd6cdd3f4e097a831fb6f5b7fd27543c7b18f Mon Sep 17 00:00:00 2001 From: nahso Date: Tue, 19 Sep 2023 15:10:28 +0800 Subject: [PATCH 04/10] pre-commit --- deepmd/descriptor/__init__.py | 6 +-- deepmd/descriptor/se_a.py | 43 +++++++++---------- deepmd/descriptor/se_a_v2.py | 11 +++-- deepmd/descriptor/se_atten.py | 15 ++++--- deepmd/utils/compress.py | 15 +++++-- deepmd/utils/tabulate.py | 20 ++++++--- ...odel_compression_se_a_ebd_type_one_side.py | 4 +- 7 files changed, 67 insertions(+), 47 deletions(-) diff --git a/deepmd/descriptor/__init__.py b/deepmd/descriptor/__init__.py index 6bb6d4a09d..4aca9c2d07 100644 --- a/deepmd/descriptor/__init__.py +++ b/deepmd/descriptor/__init__.py @@ -11,9 +11,6 @@ from .se_a import ( DescrptSeA, ) -from .se_a_v2 import ( - DescrptSeAV2, -) from .se_a_ebd import ( DescrptSeAEbd, ) @@ -24,6 +21,9 @@ from .se_a_mask import ( DescrptSeAMask, ) +from .se_a_v2 import ( + DescrptSeAV2, +) from .se_atten import ( DescrptSeAtten, ) diff --git a/deepmd/descriptor/se_a.py b/deepmd/descriptor/se_a.py index 1f75dabe49..c7e114d3b4 100644 --- a/deepmd/descriptor/se_a.py +++ b/deepmd/descriptor/se_a.py @@ -10,8 +10,8 @@ from deepmd.common import ( cast_precision, get_activation_func, - get_precision, get_np_precision, + get_precision, ) from deepmd.env import ( GLOBAL_NP_FLOAT_PRECISION, @@ -31,10 +31,17 @@ from deepmd.nvnmd.utils.config import ( nvnmd_cfg, ) +from deepmd.utils.compress import ( + get_extra_side_embedding_net_variable, + get_two_side_type_embedding, + get_type_embedding, + make_data, +) from deepmd.utils.errors import ( GraphWithoutTensorError, ) from deepmd.utils.graph import ( + get_pattern_nodes_from_graph_def, get_tensor_by_name_from_graph, ) from deepmd.utils.network import ( @@ -53,18 +60,6 @@ from deepmd.utils.type_embed import ( embed_atom_type, ) -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.compress import ( - get_type_embedding, - get_two_side_type_embedding, - get_extra_side_embedding_net_variable, - make_data, -) from .descriptor import ( Descriptor, @@ -501,14 +496,18 @@ def enable_compression( self.matrix = get_extra_side_embedding_net_variable( self, graph_def, "two_side", "matrix", suffix ) - self.bias = get_extra_side_embedding_net_variable(self, graph_def, "two_side", "bias", suffix) + self.bias = get_extra_side_embedding_net_variable( + self, graph_def, "two_side", "bias", suffix + ) self.extra_embedding = make_data(self, self.final_type_embedding) else: self.final_type_embedding = get_type_embedding(self, graph) self.matrix = get_extra_side_embedding_net_variable( self, graph_def, "one_side", "matrix", suffix ) - self.bias = get_extra_side_embedding_net_variable(self, graph_def, "one_side", "bias", suffix) + self.bias = get_extra_side_embedding_net_variable( + self, graph_def, "one_side", "bias", suffix + ) self.extra_embedding = make_data(self, self.final_type_embedding) self.compress = True @@ -983,7 +982,9 @@ def _filter_lower( initial_variables=self.extra_embeeding_net_variables, mixed_prec=self.mixed_prec, ) - net_output = tf.nn.embedding_lookup(net_output, self.nei_type_vec) + net_output = tf.nn.embedding_lookup( + net_output, self.nei_type_vec + ) else: type_embedding_nei = tf.tile( tf.reshape(type_embedding, [1, padding_ntypes, -1]), @@ -1024,9 +1025,7 @@ def _filter_lower( initial_variables=self.extra_embeeding_net_variables, mixed_prec=self.mixed_prec, ) - net_output = tf.nn.embedding_lookup( - net_output, idx - ) + net_output = tf.nn.embedding_lookup(net_output, idx) net_output = tf.reshape(net_output, [-1, self.filter_neuron[-1]]) else: xyz_scatter = self._concat_type_embedding( @@ -1077,7 +1076,7 @@ def _filter_lower( tf.reshape(inputs_i, [natom, shape_i[1] // 4, 4]), net_output, last_layer_size=outputs_size[-1], - is_sorted=False + is_sorted=False, ) else: if self.type_one_side: @@ -1333,11 +1332,11 @@ def init_variables( self.dstd = new_dstd if self.original_sel is None: self.original_sel = sel + @property def explicit_ntypes(self) -> bool: """Explicit ntypes with type embedding.""" - print('inner se_a, explicit_ntypes = ', self.stripped_type_embedding) + print("inner se_a, explicit_ntypes = ", self.stripped_type_embedding) if self.stripped_type_embedding: return True return False - diff --git a/deepmd/descriptor/se_a_v2.py b/deepmd/descriptor/se_a_v2.py index 521a63ad1a..4be27d51a5 100644 --- a/deepmd/descriptor/se_a_v2.py +++ b/deepmd/descriptor/se_a_v2.py @@ -5,24 +5,23 @@ Optional, ) +from deepmd.utils.spin import ( + Spin, +) + from .descriptor import ( Descriptor, ) from .se_a import ( DescrptSeA, ) -from deepmd.utils.spin import ( - Spin, -) log = logging.getLogger(__name__) @Descriptor.register("se_e2_a_v2") class DescrptSeAV2(DescrptSeA): - r""" - A warpper for DescriptorSeA, which set stripped_type_embedding=True - """ + r"""A warpper for DescriptorSeA, which set stripped_type_embedding=True.""" def __init__( self, diff --git a/deepmd/descriptor/se_atten.py b/deepmd/descriptor/se_atten.py index 5c8d40f95e..b0c65108e5 100644 --- a/deepmd/descriptor/se_atten.py +++ b/deepmd/descriptor/se_atten.py @@ -35,6 +35,11 @@ from deepmd.nvnmd.utils.config import ( nvnmd_cfg, ) +from deepmd.utils.compress import ( + get_extra_side_embedding_net_variable, + get_two_side_type_embedding, + make_data, +) from deepmd.utils.graph import ( get_attention_layer_variables_from_graph_def, get_pattern_nodes_from_graph_def, @@ -51,12 +56,6 @@ from deepmd.utils.tabulate import ( DPTabulate, ) -from deepmd.utils.compress import ( - get_type_embedding, - get_two_side_type_embedding, - get_extra_side_embedding_net_variable, - make_data, -) from .descriptor import ( Descriptor, @@ -425,7 +424,9 @@ def enable_compression( self.matrix = get_extra_side_embedding_net_variable( self, graph_def, "two_side", "matrix", suffix ) - self.bias = get_extra_side_embedding_net_variable(self, graph_def, "two_side", "bias", suffix) + self.bias = get_extra_side_embedding_net_variable( + self, graph_def, "two_side", "bias", suffix + ) self.two_embd = make_data(self, self.final_type_embedding) self.davg = get_tensor_by_name_from_graph( diff --git a/deepmd/utils/compress.py b/deepmd/utils/compress.py index 7e394ed897..c6e68dfe19 100644 --- a/deepmd/utils/compress.py +++ b/deepmd/utils/compress.py @@ -1,22 +1,25 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later import logging +import numpy as np + from deepmd.env import ( tf, ) - from deepmd.utils.graph import ( get_pattern_nodes_from_graph_def, get_tensor_by_name_from_graph, ) -import numpy as np log = logging.getLogger(__name__) + def get_type_embedding(self, graph): type_embedding = get_tensor_by_name_from_graph(graph, "t_typeebd") type_embedding = type_embedding.astype(self.filter_np_precision) return type_embedding + 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) @@ -38,7 +41,10 @@ def get_two_side_type_embedding(self, graph): ) return two_side_type_embedding -def get_extra_side_embedding_net_variable(self, graph_def, type_side, varialbe_name, suffix): + +def get_extra_side_embedding_net_variable( + self, graph_def, type_side, varialbe_name, suffix +): ret = {} for i in range(1, self.layer_size + 1): target = get_pattern_nodes_from_graph_def( @@ -49,13 +55,16 @@ def get_extra_side_embedding_net_variable(self, graph_def, type_side, varialbe_n 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): diff --git a/deepmd/utils/tabulate.py b/deepmd/utils/tabulate.py index 81eacaa459..546bf1b7e7 100644 --- a/deepmd/utils/tabulate.py +++ b/deepmd/utils/tabulate.py @@ -176,7 +176,9 @@ 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.DescrptSeAtten) or isinstance(self.descrpt, deepmd.descriptor.DescrptSeAV2): + if isinstance(self.descrpt, deepmd.descriptor.DescrptSeAtten) or isinstance( + self.descrpt, deepmd.descriptor.DescrptSeAV2 + ): uu = np.max(upper) ll = np.min(lower) xx = np.arange(ll, uu, stride0, dtype=self.data_type) @@ -419,7 +421,9 @@ def _get_bias(self): bias = {} for layer in range(1, self.layer_size + 1): bias["layer_" + str(layer)] = [] - if isinstance(self.descrpt, deepmd.descriptor.DescrptSeAtten) or isinstance(self.descrpt, deepmd.descriptor.DescrptSeAV2): + if isinstance(self.descrpt, deepmd.descriptor.DescrptSeAtten) or isinstance( + self.descrpt, deepmd.descriptor.DescrptSeAV2 + ): node = self.embedding_net_nodes[ f"filter_type_all{self.suffix}/bias_{layer}" ] @@ -483,7 +487,9 @@ def _get_matrix(self): matrix = {} for layer in range(1, self.layer_size + 1): matrix["layer_" + str(layer)] = [] - if isinstance(self.descrpt, deepmd.descriptor.DescrptSeAtten) or isinstance(self.descrpt, deepmd.descriptor.DescrptSeAV2): + if isinstance(self.descrpt, deepmd.descriptor.DescrptSeAtten) or isinstance( + self.descrpt, deepmd.descriptor.DescrptSeAV2 + ): node = self.embedding_net_nodes[ f"filter_type_all{self.suffix}/matrix_{layer}" ] @@ -687,7 +693,9 @@ def _spline5_switch(self, xx, rmin, rmax): def _get_layer_size(self): layer_size = 0 - if isinstance(self.descrpt, deepmd.descriptor.DescrptSeAtten) or isinstance(self.descrpt, deepmd.descriptor.DescrptSeAV2): + if isinstance(self.descrpt, deepmd.descriptor.DescrptSeAtten) or isinstance( + self.descrpt, deepmd.descriptor.DescrptSeAV2 + ): layer_size = len(self.embedding_net_nodes) // 2 elif isinstance(self.descrpt, deepmd.descriptor.DescrptSeA): layer_size = len(self.embedding_net_nodes) // ( @@ -737,7 +745,9 @@ def _all_excluded(self, ii: int) -> bool: def _get_table_size(self): table_size = 0 - if isinstance(self.descrpt, deepmd.descriptor.DescrptSeAtten) or isinstance(self.descrpt, deepmd.descriptor.DescrptSeAV2): + if isinstance(self.descrpt, deepmd.descriptor.DescrptSeAtten) or isinstance( + self.descrpt, deepmd.descriptor.DescrptSeAV2 + ): table_size = 1 elif isinstance(self.descrpt, deepmd.descriptor.DescrptSeA): table_size = self.ntypes * self.ntypes diff --git a/source/tests/test_model_compression_se_a_ebd_type_one_side.py b/source/tests/test_model_compression_se_a_ebd_type_one_side.py index 9787a4e68a..4a7c29f8ab 100644 --- a/source/tests/test_model_compression_se_a_ebd_type_one_side.py +++ b/source/tests/test_model_compression_se_a_ebd_type_one_side.py @@ -72,7 +72,9 @@ def _init_models(): def _init_models_exclude_types(): data_file = str(tests_path / os.path.join("model_compression", "data")) frozen_model = str(tests_path / "dp-original-exclude-types-se-e2-a-v2-one-side.pb") - compressed_model = str(tests_path / "dp-compressed-exclude-types-se-e2-a-v2-one-side.pb") + compressed_model = str( + tests_path / "dp-compressed-exclude-types-se-e2-a-v2-one-side.pb" + ) INPUT = str(tests_path / "input.json") jdata = j_loader(str(tests_path / os.path.join("model_compression", "input.json"))) jdata["model"]["descriptor"] = {} From 04bf0afacd4ef6d38618edd7a7cd8b010a595f16 Mon Sep 17 00:00:00 2001 From: nahso Date: Thu, 21 Sep 2023 10:46:24 +0800 Subject: [PATCH 05/10] remove debug message --- deepmd/descriptor/se_a.py | 1 - 1 file changed, 1 deletion(-) diff --git a/deepmd/descriptor/se_a.py b/deepmd/descriptor/se_a.py index c7e114d3b4..d174e26397 100644 --- a/deepmd/descriptor/se_a.py +++ b/deepmd/descriptor/se_a.py @@ -1336,7 +1336,6 @@ def init_variables( @property def explicit_ntypes(self) -> bool: """Explicit ntypes with type embedding.""" - print("inner se_a, explicit_ntypes = ", self.stripped_type_embedding) if self.stripped_type_embedding: return True return False From e55778c3b425a0edbea4cb000166d0cc9d83b26c Mon Sep 17 00:00:00 2001 From: nahso Date: Mon, 25 Sep 2023 13:53:18 +0800 Subject: [PATCH 06/10] rename to se_a_ebd_v2, add model ut --- deepmd/descriptor/__init__.py | 6 +- deepmd/descriptor/se_a.py | 4 + .../descriptor/{se_a_v2.py => se_a_ebd_v2.py} | 10 +- deepmd/utils/argcheck.py | 4 +- deepmd/utils/tabulate.py | 10 +- .../se_e2_a_tebd_compressible/input.json | 98 +++++++++++++++++++ .../tests/test_model_compression_se_a_ebd.py | 12 +-- ...odel_compression_se_a_ebd_type_one_side.py | 4 +- 8 files changed, 127 insertions(+), 21 deletions(-) rename deepmd/descriptor/{se_a_v2.py => se_a_ebd_v2.py} (86%) create mode 100644 examples/water/se_e2_a_tebd_compressible/input.json diff --git a/deepmd/descriptor/__init__.py b/deepmd/descriptor/__init__.py index 4aca9c2d07..fa35e05888 100644 --- a/deepmd/descriptor/__init__.py +++ b/deepmd/descriptor/__init__.py @@ -21,8 +21,8 @@ from .se_a_mask import ( DescrptSeAMask, ) -from .se_a_v2 import ( - DescrptSeAV2, +from .se_a_ebd_v2 import ( + DescrptSeAEbdV2, ) from .se_atten import ( DescrptSeAtten, @@ -42,7 +42,7 @@ "DescrptHybrid", "DescrptLocFrame", "DescrptSeA", - "DescrptSeAV2", + "DescrptSeAEbdV2", "DescrptSeAEbd", "DescrptSeAEf", "DescrptSeAEfLower", diff --git a/deepmd/descriptor/se_a.py b/deepmd/descriptor/se_a.py index d174e26397..7f2db3bc83 100644 --- a/deepmd/descriptor/se_a.py +++ b/deepmd/descriptor/se_a.py @@ -743,6 +743,10 @@ def _pass_filter( type_embedding = input_dict.get("type_embedding", None) else: type_embedding = None + if self.stripped_type_embedding and type_embedding is None: + raise RuntimeError( + "type_embedding is required for se_a_tebd_v2 model." + ) start_index = 0 inputs = tf.reshape(inputs, [-1, natoms[0], self.ndescrpt]) output = [] diff --git a/deepmd/descriptor/se_a_v2.py b/deepmd/descriptor/se_a_ebd_v2.py similarity index 86% rename from deepmd/descriptor/se_a_v2.py rename to deepmd/descriptor/se_a_ebd_v2.py index 4be27d51a5..c6e3cebc71 100644 --- a/deepmd/descriptor/se_a_v2.py +++ b/deepmd/descriptor/se_a_ebd_v2.py @@ -19,9 +19,13 @@ log = logging.getLogger(__name__) -@Descriptor.register("se_e2_a_v2") -class DescrptSeAV2(DescrptSeA): - r"""A warpper for DescriptorSeA, which set stripped_type_embedding=True.""" +@Descriptor.register("se_a_tpe_v2") +@Descriptor.register("se_a_ebd_v2") +class DescrptSeAEbdV2(DescrptSeA): + r"""A compressible se_a_ebd model. + + This model is a warpper for DescriptorSeA, which set stripped_type_embedding=True. + """ def __init__( self, diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 4652404e9c..ae446ef348 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -432,8 +432,8 @@ def descrpt_se_atten_v2_args(): ] -@descrpt_args_plugin.register("se_e2_a_v2") -def descrpt_se_e2_a_v2_args(): +@descrpt_args_plugin.register("se_a_ebd_v2", alias=["se_a_tpe_v2"]) +def descrpt_se_a_ebd_v2_args(): return descrpt_se_a_args() diff --git a/deepmd/utils/tabulate.py b/deepmd/utils/tabulate.py index 546bf1b7e7..427887089a 100644 --- a/deepmd/utils/tabulate.py +++ b/deepmd/utils/tabulate.py @@ -177,7 +177,7 @@ 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.DescrptSeAtten) or isinstance( - self.descrpt, deepmd.descriptor.DescrptSeAV2 + self.descrpt, deepmd.descriptor.DescrptSeAEbdV2 ): uu = np.max(upper) ll = np.min(lower) @@ -422,7 +422,7 @@ def _get_bias(self): for layer in range(1, self.layer_size + 1): bias["layer_" + str(layer)] = [] if isinstance(self.descrpt, deepmd.descriptor.DescrptSeAtten) or isinstance( - self.descrpt, deepmd.descriptor.DescrptSeAV2 + self.descrpt, deepmd.descriptor.DescrptSeAEbdV2 ): node = self.embedding_net_nodes[ f"filter_type_all{self.suffix}/bias_{layer}" @@ -488,7 +488,7 @@ def _get_matrix(self): for layer in range(1, self.layer_size + 1): matrix["layer_" + str(layer)] = [] if isinstance(self.descrpt, deepmd.descriptor.DescrptSeAtten) or isinstance( - self.descrpt, deepmd.descriptor.DescrptSeAV2 + self.descrpt, deepmd.descriptor.DescrptSeAEbdV2 ): node = self.embedding_net_nodes[ f"filter_type_all{self.suffix}/matrix_{layer}" @@ -694,7 +694,7 @@ def _spline5_switch(self, xx, rmin, rmax): def _get_layer_size(self): layer_size = 0 if isinstance(self.descrpt, deepmd.descriptor.DescrptSeAtten) or isinstance( - self.descrpt, deepmd.descriptor.DescrptSeAV2 + self.descrpt, deepmd.descriptor.DescrptSeAEbdV2 ): layer_size = len(self.embedding_net_nodes) // 2 elif isinstance(self.descrpt, deepmd.descriptor.DescrptSeA): @@ -746,7 +746,7 @@ def _all_excluded(self, ii: int) -> bool: def _get_table_size(self): table_size = 0 if isinstance(self.descrpt, deepmd.descriptor.DescrptSeAtten) or isinstance( - self.descrpt, deepmd.descriptor.DescrptSeAV2 + self.descrpt, deepmd.descriptor.DescrptSeAEbdV2 ): table_size = 1 elif isinstance(self.descrpt, deepmd.descriptor.DescrptSeA): diff --git a/examples/water/se_e2_a_tebd_compressible/input.json b/examples/water/se_e2_a_tebd_compressible/input.json new file mode 100644 index 0000000000..1ec85a968c --- /dev/null +++ b/examples/water/se_e2_a_tebd_compressible/input.json @@ -0,0 +1,98 @@ +{ + "_comment1": " model parameters", + "model": { + "type_map": [ + "O", + "H" + ], + "type_embedding": { + "neuron": [ + 2, + 4, + 8 + ], + "resnet_dt": false, + "precision": "float64", + "seed": 1 + }, + "descriptor": { + "type": "se_a_ebd_v2", + "sel": [ + 46, + 92 + ], + "rcut_smth": 0.50, + "rcut": 6.00, + "neuron": [ + 25, + 50, + 100 + ], + "resnet_dt": false, + "axis_neuron": 16, + "type_one_side": true, + "precision": "float64", + "seed": 1, + "_comment2": " that's all" + }, + "fitting_net": { + "neuron": [ + 240, + 240, + 240 + ], + "resnet_dt": true, + "precision": "float64", + "seed": 1, + "_comment3": " that's all" + }, + "_comment4": " that's all" + }, + + "learning_rate": { + "type": "exp", + "decay_steps": 5000, + "start_lr": 0.001, + "stop_lr": 3.51e-8, + "_comment5": "that's all" + }, + + "loss": { + "type": "ener", + "start_pref_e": 0.02, + "limit_pref_e": 1, + "start_pref_f": 1000, + "limit_pref_f": 1, + "start_pref_v": 0, + "limit_pref_v": 0, + "_comment6": " that's all" + }, + + "training": { + "training_data": { + "systems": [ + "../data/data_0/", + "../data/data_1/", + "../data/data_2/" + ], + "batch_size": "auto", + "_comment7": "that's all" + }, + "validation_data": { + "systems": [ + "../data/data_3" + ], + "batch_size": 1, + "numb_btch": 3, + "_comment8": "that's all" + }, + "numb_steps": 1000000, + "seed": 10, + "disp_file": "lcurve.out", + "disp_freq": 100, + "save_freq": 1000, + "_comment9": "that's all" + }, + + "_comment10": "that's all" +} diff --git a/source/tests/test_model_compression_se_a_ebd.py b/source/tests/test_model_compression_se_a_ebd.py index 78ef639a11..2a3163b062 100644 --- a/source/tests/test_model_compression_se_a_ebd.py +++ b/source/tests/test_model_compression_se_a_ebd.py @@ -48,12 +48,12 @@ def _init_models(): data_file = str(tests_path / os.path.join("model_compression", "data")) frozen_model = str(tests_path / "dp-original-se-e2-a-v2.pb") compressed_model = str(tests_path / "dp-compressed-se-e2-a-v2.pb") - INPUT = str(tests_path / "input_se_e2_a_v2.json") + INPUT = str(tests_path / "input_se_a_ebd_v2.json") jdata = j_loader(str(tests_path / os.path.join("model_compression", "input.json"))) jdata["training"]["training_data"]["systems"] = data_file jdata["training"]["validation_data"]["systems"] = data_file jdata["model"]["descriptor"] = {} - jdata["model"]["descriptor"]["type"] = "se_e2_a_v2" + jdata["model"]["descriptor"]["type"] = "se_a_ebd_v2" jdata["model"]["descriptor"]["type_one_side"] = False jdata["model"]["type_embedding"] = {} jdata["model"]["type_embedding"]["neuron"] = [1] @@ -73,12 +73,12 @@ def _init_models(): def _init_models_exclude_types(): data_file = str(tests_path / os.path.join("model_compression", "data")) - frozen_model = str(tests_path / "dp-original-exclude-types-se-e2-a-v2.pb") - compressed_model = str(tests_path / "dp-compressed-exclude-types-se-e2-a-v2.pb") - INPUT = str(tests_path / "input_se_e2_a_v2.json") + frozen_model = str(tests_path / "dp-original-exclude-types-se-e2-a-ebd-v2.pb") + compressed_model = str(tests_path / "dp-compressed-exclude-types-se-e2-a-ebd-v2.pb") + INPUT = str(tests_path / "input_se_a_ebd_v2.json") jdata = j_loader(str(tests_path / os.path.join("model_compression", "input.json"))) jdata["model"]["descriptor"] = {} - jdata["model"]["descriptor"]["type"] = "se_e2_a_v2" + jdata["model"]["descriptor"]["type"] = "se_a_ebd_v2" jdata["model"]["descriptor"]["exclude_types"] = [[0, 1]] jdata["model"]["descriptor"]["type_one_side"] = False jdata["model"]["type_embedding"] = {} diff --git a/source/tests/test_model_compression_se_a_ebd_type_one_side.py b/source/tests/test_model_compression_se_a_ebd_type_one_side.py index 4a7c29f8ab..2f3d16b05f 100644 --- a/source/tests/test_model_compression_se_a_ebd_type_one_side.py +++ b/source/tests/test_model_compression_se_a_ebd_type_one_side.py @@ -53,7 +53,7 @@ def _init_models(): jdata["training"]["training_data"]["systems"] = data_file jdata["training"]["validation_data"]["systems"] = data_file jdata["model"]["descriptor"] = {} - jdata["model"]["descriptor"]["type"] = "se_e2_a_v2" + jdata["model"]["descriptor"]["type"] = "se_a_ebd_v2" jdata["model"]["descriptor"]["type_one_side"] = True jdata["model"]["type_embedding"] = {} jdata["model"]["type_embedding"]["neuron"] = [1] @@ -78,7 +78,7 @@ def _init_models_exclude_types(): INPUT = str(tests_path / "input.json") jdata = j_loader(str(tests_path / os.path.join("model_compression", "input.json"))) jdata["model"]["descriptor"] = {} - jdata["model"]["descriptor"]["type"] = "se_e2_a_v2" + jdata["model"]["descriptor"]["type"] = "se_a_ebd_v2" jdata["model"]["descriptor"]["exclude_types"] = [[0, 1]] jdata["model"]["descriptor"]["type_one_side"] = True jdata["model"]["type_embedding"] = {} From dce0618cc40bbfeef15733ae1f978be945082f51 Mon Sep 17 00:00:00 2001 From: nahso Date: Mon, 25 Sep 2023 13:57:24 +0800 Subject: [PATCH 07/10] pre-commit --- deepmd/descriptor/__init__.py | 6 +++--- deepmd/descriptor/se_a.py | 4 +--- 2 files changed, 4 insertions(+), 6 deletions(-) diff --git a/deepmd/descriptor/__init__.py b/deepmd/descriptor/__init__.py index fa35e05888..6ef6c50da5 100644 --- a/deepmd/descriptor/__init__.py +++ b/deepmd/descriptor/__init__.py @@ -14,6 +14,9 @@ from .se_a_ebd import ( DescrptSeAEbd, ) +from .se_a_ebd_v2 import ( + DescrptSeAEbdV2, +) from .se_a_ef import ( DescrptSeAEf, DescrptSeAEfLower, @@ -21,9 +24,6 @@ from .se_a_mask import ( DescrptSeAMask, ) -from .se_a_ebd_v2 import ( - DescrptSeAEbdV2, -) from .se_atten import ( DescrptSeAtten, ) diff --git a/deepmd/descriptor/se_a.py b/deepmd/descriptor/se_a.py index 7f2db3bc83..8f0051cd4e 100644 --- a/deepmd/descriptor/se_a.py +++ b/deepmd/descriptor/se_a.py @@ -744,9 +744,7 @@ def _pass_filter( else: type_embedding = None if self.stripped_type_embedding and type_embedding is None: - raise RuntimeError( - "type_embedding is required for se_a_tebd_v2 model." - ) + raise RuntimeError("type_embedding is required for se_a_tebd_v2 model.") start_index = 0 inputs = tf.reshape(inputs, [-1, natoms[0], self.ndescrpt]) output = [] From 925f351a3157d32e05778156b7a0d8d490d71cf0 Mon Sep 17 00:00:00 2001 From: nahso Date: Mon, 25 Sep 2023 13:59:37 +0800 Subject: [PATCH 08/10] ut for se_e2_a_ebd_v2 model --- source/tests/test_model_se_a_ebd_v2.py | 159 +++++++++++++++++++++++++ 1 file changed, 159 insertions(+) create mode 100644 source/tests/test_model_se_a_ebd_v2.py diff --git a/source/tests/test_model_se_a_ebd_v2.py b/source/tests/test_model_se_a_ebd_v2.py new file mode 100644 index 0000000000..e91a0dbf1f --- /dev/null +++ b/source/tests/test_model_se_a_ebd_v2.py @@ -0,0 +1,159 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import numpy as np +from common import ( + DataSystem, + gen_data, + j_loader, +) + +from deepmd.common import ( + j_must_have, +) +from deepmd.descriptor.se_a_ebd_v2 import ( + DescrptSeAEbdV2, +) +from deepmd.env import ( + tf, +) +from deepmd.fit import ( + EnerFitting, +) +from deepmd.model import ( + EnerModel, +) +from deepmd.utils.type_embed import ( + TypeEmbedNet, +) + +GLOBAL_ENER_FLOAT_PRECISION = tf.float64 +GLOBAL_TF_FLOAT_PRECISION = tf.float64 +GLOBAL_NP_FLOAT_PRECISION = np.float64 + + +class TestModel(tf.test.TestCase): + def setUp(self): + gen_data() + + def test_model(self): + jfile = "water_se_a_ebd.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"]["type_embedding"] = {} + jdata["model"]["type_embedding"]["neuron"] = [1] + jdata["model"]["type_embedding"]["resnet_dt"] = False + jdata["model"]["type_embedding"]["seed"] = 1 + 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, + ) + descrpt = DescrptSeAEbdV2( + **jdata["model"]["descriptor"], + ) + jdata["model"]["fitting_net"]["descrpt"] = descrpt + fitting = EnerFitting( + **jdata["model"]["fitting_net"], + ) + # fitting = EnerFitting(jdata['model']['fitting_net'], descrpt) + 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 + + model_pred = model.build( + t_coord, + t_type, + t_natoms, + t_box, + t_mesh, + t_fparam, + suffix="se_a_ebd_v2", + 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.cached_session().__enter__() + sess.run(tf.global_variables_initializer()) + [e, f, v] = sess.run([energy, force, virial], feed_dict=feed_dict_test) + + e = e.reshape([-1]) + f = f.reshape([-1]) + v = v.reshape([-1]) + + refe = [5.435394596262052014e-01 ] + reff = [ + 6.583728125594628944e-02,7.228993116083935744e-02,1.971543579114074483e-03,6.567474563776359853e-02,7.809421727465599983e-02,-4.866958849094786890e-03,-8.670511901715304004e-02,3.525374157021862048e-02,1.415748959800727487e-03,6.375813001810648473e-02,-1.139053242798149790e-01,-4.178593754384440744e-03,-1.471737787218250215e-01,4.189712704724830872e-02,7.011731363309440038e-03,3.860874082716164030e-02,-1.136296927731473005e-01,-1.353471298745012206e-03 + ] + refv = [ + -4.243979601186427253e-01,1.097173849143971286e-01,1.227299373463585502e-02,1.097173849143970314e-01,-2.462891443164323124e-01,-5.711664180530139426e-03,1.227299373463585502e-02,-5.711664180530143763e-03,-6.217348853341628408e-04 + ] + refe = np.reshape(refe, [-1]) + reff = np.reshape(reff, [-1]) + refv = np.reshape(refv, [-1]) + + places = 6 + for ii in range(e.size): + self.assertAlmostEqual(e[ii], refe[ii], places=places) + for ii in range(f.size): + self.assertAlmostEqual(f[ii], reff[ii], places=places) + for ii in range(v.size): + self.assertAlmostEqual(v[ii], refv[ii], places=places) From 8bb548186f4724a8d6977a61e6f325b5b337a336 Mon Sep 17 00:00:00 2001 From: nahso Date: Mon, 25 Sep 2023 14:00:25 +0800 Subject: [PATCH 09/10] pre-commit --- source/tests/test_model_se_a_ebd_v2.py | 31 +++++++++++++++++++++++--- 1 file changed, 28 insertions(+), 3 deletions(-) diff --git a/source/tests/test_model_se_a_ebd_v2.py b/source/tests/test_model_se_a_ebd_v2.py index e91a0dbf1f..71860890ce 100644 --- a/source/tests/test_model_se_a_ebd_v2.py +++ b/source/tests/test_model_se_a_ebd_v2.py @@ -139,12 +139,37 @@ def test_model(self): f = f.reshape([-1]) v = v.reshape([-1]) - refe = [5.435394596262052014e-01 ] + refe = [5.435394596262052014e-01] reff = [ - 6.583728125594628944e-02,7.228993116083935744e-02,1.971543579114074483e-03,6.567474563776359853e-02,7.809421727465599983e-02,-4.866958849094786890e-03,-8.670511901715304004e-02,3.525374157021862048e-02,1.415748959800727487e-03,6.375813001810648473e-02,-1.139053242798149790e-01,-4.178593754384440744e-03,-1.471737787218250215e-01,4.189712704724830872e-02,7.011731363309440038e-03,3.860874082716164030e-02,-1.136296927731473005e-01,-1.353471298745012206e-03 + 6.583728125594628944e-02, + 7.228993116083935744e-02, + 1.971543579114074483e-03, + 6.567474563776359853e-02, + 7.809421727465599983e-02, + -4.866958849094786890e-03, + -8.670511901715304004e-02, + 3.525374157021862048e-02, + 1.415748959800727487e-03, + 6.375813001810648473e-02, + -1.139053242798149790e-01, + -4.178593754384440744e-03, + -1.471737787218250215e-01, + 4.189712704724830872e-02, + 7.011731363309440038e-03, + 3.860874082716164030e-02, + -1.136296927731473005e-01, + -1.353471298745012206e-03, ] refv = [ - -4.243979601186427253e-01,1.097173849143971286e-01,1.227299373463585502e-02,1.097173849143970314e-01,-2.462891443164323124e-01,-5.711664180530139426e-03,1.227299373463585502e-02,-5.711664180530143763e-03,-6.217348853341628408e-04 + -4.243979601186427253e-01, + 1.097173849143971286e-01, + 1.227299373463585502e-02, + 1.097173849143970314e-01, + -2.462891443164323124e-01, + -5.711664180530139426e-03, + 1.227299373463585502e-02, + -5.711664180530143763e-03, + -6.217348853341628408e-04, ] refe = np.reshape(refe, [-1]) reff = np.reshape(reff, [-1]) From 487e6a388b90f5373e787813a0d1233ed530bb3e Mon Sep 17 00:00:00 2001 From: nahso Date: Mon, 25 Sep 2023 15:34:50 +0800 Subject: [PATCH 10/10] remove old se_e2_a_tebd/input.json --- examples/water/se_e2_a_tebd/input.json | 2 +- .../se_e2_a_tebd_compressible/input.json | 98 ------------------- 2 files changed, 1 insertion(+), 99 deletions(-) delete mode 100644 examples/water/se_e2_a_tebd_compressible/input.json diff --git a/examples/water/se_e2_a_tebd/input.json b/examples/water/se_e2_a_tebd/input.json index 101c1a7a4f..1ec85a968c 100644 --- a/examples/water/se_e2_a_tebd/input.json +++ b/examples/water/se_e2_a_tebd/input.json @@ -16,7 +16,7 @@ "seed": 1 }, "descriptor": { - "type": "se_e2_a", + "type": "se_a_ebd_v2", "sel": [ 46, 92 diff --git a/examples/water/se_e2_a_tebd_compressible/input.json b/examples/water/se_e2_a_tebd_compressible/input.json deleted file mode 100644 index 1ec85a968c..0000000000 --- a/examples/water/se_e2_a_tebd_compressible/input.json +++ /dev/null @@ -1,98 +0,0 @@ -{ - "_comment1": " model parameters", - "model": { - "type_map": [ - "O", - "H" - ], - "type_embedding": { - "neuron": [ - 2, - 4, - 8 - ], - "resnet_dt": false, - "precision": "float64", - "seed": 1 - }, - "descriptor": { - "type": "se_a_ebd_v2", - "sel": [ - 46, - 92 - ], - "rcut_smth": 0.50, - "rcut": 6.00, - "neuron": [ - 25, - 50, - 100 - ], - "resnet_dt": false, - "axis_neuron": 16, - "type_one_side": true, - "precision": "float64", - "seed": 1, - "_comment2": " that's all" - }, - "fitting_net": { - "neuron": [ - 240, - 240, - 240 - ], - "resnet_dt": true, - "precision": "float64", - "seed": 1, - "_comment3": " that's all" - }, - "_comment4": " that's all" - }, - - "learning_rate": { - "type": "exp", - "decay_steps": 5000, - "start_lr": 0.001, - "stop_lr": 3.51e-8, - "_comment5": "that's all" - }, - - "loss": { - "type": "ener", - "start_pref_e": 0.02, - "limit_pref_e": 1, - "start_pref_f": 1000, - "limit_pref_f": 1, - "start_pref_v": 0, - "limit_pref_v": 0, - "_comment6": " that's all" - }, - - "training": { - "training_data": { - "systems": [ - "../data/data_0/", - "../data/data_1/", - "../data/data_2/" - ], - "batch_size": "auto", - "_comment7": "that's all" - }, - "validation_data": { - "systems": [ - "../data/data_3" - ], - "batch_size": 1, - "numb_btch": 3, - "_comment8": "that's all" - }, - "numb_steps": 1000000, - "seed": 10, - "disp_file": "lcurve.out", - "disp_freq": 100, - "save_freq": 1000, - "_comment9": "that's all" - }, - - "_comment10": "that's all" -}