From 9b7aec3a177d022302ae564952899668159a6257 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Mon, 13 Sep 2021 19:01:09 -0400 Subject: [PATCH 1/5] support init-frz-model for se_r and se_t I just copied several codes from se_a, which can also be applied to se_r and se_t. --- deepmd/descriptor/se_r.py | 70 ++++++++++++++++++++++++++++++++++++--- deepmd/descriptor/se_t.py | 69 +++++++++++++++++++++++++++++++++++--- 2 files changed, 129 insertions(+), 10 deletions(-) diff --git a/deepmd/descriptor/se_r.py b/deepmd/descriptor/se_r.py index 8b014a2501..41abc4c677 100644 --- a/deepmd/descriptor/se_r.py +++ b/deepmd/descriptor/se_r.py @@ -9,6 +9,7 @@ from deepmd.env import op_module from deepmd.env import default_tf_session_config from deepmd.utils.network import embedding_net, embedding_net_rand_seed_shift +from deepmd.utils.graph import get_embedding_net_variables from deepmd.utils.sess import run_sess from .descriptor import Descriptor @@ -114,6 +115,7 @@ def __init__ (self, self.useBN = False self.davg = None self.dstd = None + self.embedding_net_variables = None self.place_holders = {} avg_zero = np.zeros([self.ntypes,self.ndescrpt]).astype(GLOBAL_NP_FLOAT_PRECISION) @@ -312,10 +314,10 @@ def build (self, sel = self.sel_r) self.descrpt_reshape = tf.reshape(self.descrpt, [-1, self.ndescrpt]) - self.descrpt_reshape = tf.identity(self.descrpt_reshape, name = 'o_rmat') - self.descrpt_deriv = tf.identity(self.descrpt_deriv, name = 'o_rmat_deriv') - self.rij = tf.identity(self.rij, name = 'o_rij') - self.nlist = tf.identity(self.nlist, name = 'o_nlist') + self.descrpt_reshape = tf.identity(self.descrpt_reshape, name = 'o_rmat' + suffix) + self.descrpt_deriv = tf.identity(self.descrpt_deriv, name = 'o_rmat_deriv' + suffix) + self.rij = tf.identity(self.rij, name = 'o_rij' + suffix) + self.nlist = tf.identity(self.nlist, name = 'o_nlist' + suffix) # only used when tensorboard was set as true tf.summary.histogram('descrpt', self.descrpt) @@ -485,7 +487,9 @@ def _filter_r(self, bavg = bavg, seed = self.seed, trainable = trainable, - uniform_seed = self.uniform_seed) + uniform_seed = self.uniform_seed, + initial_variables = self.embedding_net_variables, + ) if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift # natom x nei_type_i x out_size xyz_scatter = tf.reshape(xyz_scatter, (-1, shape_i[1], outputs_size[-1])) @@ -502,3 +506,59 @@ def _filter_r(self, result = tf.reduce_mean(xyz_scatter, axis = 1) * res_rescale return result + + def get_tensor_names(self, suffix : str = "") -> Tuple[str]: + """Get names of tensors. + + Parameters + ---------- + suffix : str + The suffix of the scope + + Returns + ------- + Tuple[str] + Names of tensors + """ + return (f'o_rmat{suffix}:0', f'o_rmat_deriv{suffix}:0', f'o_rij{suffix}:0', f'o_nlist{suffix}:0') + + def pass_tensors_from_frz_model(self, + descrpt_reshape : tf.Tensor, + descrpt_deriv : tf.Tensor, + rij : tf.Tensor, + nlist : tf.Tensor + ): + """ + Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def + + Parameters + ---------- + descrpt_reshape + The passed descrpt_reshape tensor + descrpt_deriv + The passed descrpt_deriv tensor + rij + The passed rij tensor + nlist + The passed nlist tensor + """ + self.rij = rij + self.nlist = nlist + self.descrpt_deriv = descrpt_deriv + self.descrpt_reshape = descrpt_reshape + + def init_variables(self, + model_file : str, + suffix : str = "", + ) -> None: + """ + Init the embedding net variables with the given dict + + Parameters + ---------- + model_file : str + The input frozen model file + suffix : str, optional + The suffix of the scope + """ + self.embedding_net_variables = get_embedding_net_variables(model_file, suffix = suffix) diff --git a/deepmd/descriptor/se_t.py b/deepmd/descriptor/se_t.py index 5663a33644..e2bfb3449c 100644 --- a/deepmd/descriptor/se_t.py +++ b/deepmd/descriptor/se_t.py @@ -9,6 +9,7 @@ from deepmd.env import op_module from deepmd.env import default_tf_session_config from deepmd.utils.network import embedding_net, embedding_net_rand_seed_shift +from deepmd.utils.graph import get_embedding_net_variables from deepmd.utils.sess import run_sess from .descriptor import Descriptor @@ -97,6 +98,7 @@ def __init__ (self, self.useBN = False self.dstd = None self.davg = None + self.embedding_net_variables = None self.place_holders = {} avg_zero = np.zeros([self.ntypes,self.ndescrpt]).astype(GLOBAL_NP_FLOAT_PRECISION) @@ -311,10 +313,10 @@ def build (self, sel_r = self.sel_r) self.descrpt_reshape = tf.reshape(self.descrpt, [-1, self.ndescrpt]) - self.descrpt_reshape = tf.identity(self.descrpt_reshape, name = 'o_rmat') - self.descrpt_deriv = tf.identity(self.descrpt_deriv, name = 'o_rmat_deriv') - self.rij = tf.identity(self.rij, name = 'o_rij') - self.nlist = tf.identity(self.nlist, name = 'o_nlist') + self.descrpt_reshape = tf.identity(self.descrpt_reshape, name = 'o_rmat' + suffix) + self.descrpt_deriv = tf.identity(self.descrpt_deriv, name = 'o_rmat_deriv' + suffix) + self.rij = tf.identity(self.rij, name = 'o_rij' + suffix) + self.nlist = tf.identity(self.nlist, name = 'o_nlist' + suffix) self.dout, self.qmat = self._pass_filter(self.descrpt_reshape, atype, @@ -509,7 +511,9 @@ def _filter(self, bavg = bavg, seed = self.seed, trainable = trainable, - uniform_seed = self.uniform_seed) + uniform_seed = self.uniform_seed, + initial_variables = self.embedding_net_variables, + ) if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift # with natom x nei_type_i x nei_type_j x out_size ebd_env_ij = tf.reshape(ebd_env_ij, [-1, nei_type_i, nei_type_j, outputs_size[-1]]) @@ -522,3 +526,58 @@ def _filter(self, result += res_ij return result, None + def get_tensor_names(self, suffix : str = "") -> Tuple[str]: + """Get names of tensors. + + Parameters + ---------- + suffix : str + The suffix of the scope + + Returns + ------- + Tuple[str] + Names of tensors + """ + return (f'o_rmat{suffix}:0', f'o_rmat_deriv{suffix}:0', f'o_rij{suffix}:0', f'o_nlist{suffix}:0') + + def pass_tensors_from_frz_model(self, + descrpt_reshape : tf.Tensor, + descrpt_deriv : tf.Tensor, + rij : tf.Tensor, + nlist : tf.Tensor + ): + """ + Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def + + Parameters + ---------- + descrpt_reshape + The passed descrpt_reshape tensor + descrpt_deriv + The passed descrpt_deriv tensor + rij + The passed rij tensor + nlist + The passed nlist tensor + """ + self.rij = rij + self.nlist = nlist + self.descrpt_deriv = descrpt_deriv + self.descrpt_reshape = descrpt_reshape + + def init_variables(self, + model_file : str, + suffix : str = "", + ) -> None: + """ + Init the embedding net variables with the given dict + + Parameters + ---------- + model_file : str + The input frozen model file + suffix : str, optional + The suffix of the scope + """ + self.embedding_net_variables = get_embedding_net_variables(model_file, suffix = suffix) From b1766994b118cb744880c06331417bc642a7553e Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Mon, 13 Sep 2021 21:31:50 -0400 Subject: [PATCH 2/5] add a base class for smooth descriptors --- deepmd/descriptor/se.py | 108 ++++++++++++++++++++++++++++++++++++++ deepmd/descriptor/se_a.py | 65 ++--------------------- deepmd/descriptor/se_r.py | 65 ++--------------------- deepmd/descriptor/se_t.py | 65 ++--------------------- 4 files changed, 118 insertions(+), 185 deletions(-) create mode 100644 deepmd/descriptor/se.py diff --git a/deepmd/descriptor/se.py b/deepmd/descriptor/se.py new file mode 100644 index 0000000000..72bf908ecb --- /dev/null +++ b/deepmd/descriptor/se.py @@ -0,0 +1,108 @@ +from typing import Tuple, List + +from deepmd.env import tf +from deepmd.utils.graph import get_embedding_net_variables +from .descriptor import Descriptor + + +class DescrptSe (Descriptor): + """A base class for smooth version of descriptors. + + Notes + ----- + All of these descriptors have an environmental matrix and an + embedding network (:meth:`deepmd.utils.network.embedding_net`), so + they can share some similiar methods without defining them twice. + + Attributes + ---------- + embedding_net_variables : dict + initial embedding network variables + descrpt_reshape : tf.Tensor + the reshaped descriptor + descrpt_deriv : tf.Tensor + the descriptor derivative + rij : tf.Tensor + distances between two atoms + nlist : tf.Tensor + the neighbor list + + """ + def _identity_tensors(self, suffix : str = "") -> None: + """Identify tensors which are expected to be stored and restored. + + Notes + ----- + These tensors will be indentitied: + self.descrpt_reshape : o_rmat + self.descrpt_deriv : o_rmat_deriv + self.rij : o_rij + self.nlist : o_nlist + Thus, this method should be called during building the descriptor and + after these tensors are initialized. + + Parameters + ---------- + suffix : str + The suffix of the scope + """ + self.descrpt_reshape = tf.identity(self.descrpt_reshape, name = 'o_rmat' + suffix) + self.descrpt_deriv = tf.identity(self.descrpt_deriv, name = 'o_rmat_deriv' + suffix) + self.rij = tf.identity(self.rij, name = 'o_rij' + suffix) + self.nlist = tf.identity(self.nlist, name = 'o_nlist' + suffix) + + def get_tensor_names(self, suffix : str = "") -> Tuple[str]: + """Get names of tensors. + + Parameters + ---------- + suffix : str + The suffix of the scope + + Returns + ------- + Tuple[str] + Names of tensors + """ + return (f'o_rmat{suffix}:0', f'o_rmat_deriv{suffix}:0', f'o_rij{suffix}:0', f'o_nlist{suffix}:0') + + def pass_tensors_from_frz_model(self, + descrpt_reshape : tf.Tensor, + descrpt_deriv : tf.Tensor, + rij : tf.Tensor, + nlist : tf.Tensor + ): + """ + Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def + + Parameters + ---------- + descrpt_reshape + The passed descrpt_reshape tensor + descrpt_deriv + The passed descrpt_deriv tensor + rij + The passed rij tensor + nlist + The passed nlist tensor + """ + self.rij = rij + self.nlist = nlist + self.descrpt_deriv = descrpt_deriv + self.descrpt_reshape = descrpt_reshape + + def init_variables(self, + model_file : str, + suffix : str = "", + ) -> None: + """ + Init the embedding net variables with the given dict + + Parameters + ---------- + model_file : str + The input frozen model file + suffix : str, optional + The suffix of the scope + """ + self.embedding_net_variables = get_embedding_net_variables(model_file, suffix = suffix) diff --git a/deepmd/descriptor/se_a.py b/deepmd/descriptor/se_a.py index b3db13f019..8881ac3a14 100644 --- a/deepmd/descriptor/se_a.py +++ b/deepmd/descriptor/se_a.py @@ -15,10 +15,11 @@ from deepmd.utils.sess import run_sess from deepmd.utils.graph import load_graph_def, get_tensor_by_name_from_graph, get_embedding_net_variables from .descriptor import Descriptor +from .se import DescrptSe @Descriptor.register("se_e2_a") @Descriptor.register("se_a") -class DescrptSeA (Descriptor): +class DescrptSeA (DescrptSe): r"""DeepPot-SE constructed from all information (both angular and radial) of atomic configurations. The embedding takes the distance between atoms as input. @@ -435,10 +436,7 @@ def build (self, tf.summary.histogram('nlist', self.nlist) self.descrpt_reshape = tf.reshape(self.descrpt, [-1, self.ndescrpt]) - self.descrpt_reshape = tf.identity(self.descrpt_reshape, name = 'o_rmat' + suffix) - self.descrpt_deriv = tf.identity(self.descrpt_deriv, name = 'o_rmat_deriv' + suffix) - self.rij = tf.identity(self.rij, name = 'o_rij' + suffix) - self.nlist = tf.identity(self.nlist, name = 'o_nlist' + suffix) + self._identity_tensors(suffix=suffix) self.dout, self.qmat = self._pass_filter(self.descrpt_reshape, atype, @@ -458,63 +456,6 @@ def get_rot_mat(self) -> tf.Tensor: """ return self.qmat - def get_tensor_names(self, suffix : str = "") -> Tuple[str]: - """Get names of tensors. - - Parameters - ---------- - suffix : str - The suffix of the scope - - Returns - ------- - Tuple[str] - Names of tensors - """ - return (f'o_rmat{suffix}:0', f'o_rmat_deriv{suffix}:0', f'o_rij{suffix}:0', f'o_nlist{suffix}:0') - - def pass_tensors_from_frz_model(self, - descrpt_reshape : tf.Tensor, - descrpt_deriv : tf.Tensor, - rij : tf.Tensor, - nlist : tf.Tensor - ): - """ - Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def - - Parameters - ---------- - descrpt_reshape - The passed descrpt_reshape tensor - descrpt_deriv - The passed descrpt_deriv tensor - rij - The passed rij tensor - nlist - The passed nlist tensor - """ - self.rij = rij - self.nlist = nlist - self.descrpt_deriv = descrpt_deriv - self.descrpt_reshape = descrpt_reshape - - def init_variables(self, - model_file : str, - suffix : str = "", - ) -> None: - """ - Init the embedding net variables with the given dict - - Parameters - ---------- - model_file : str - The input frozen model file - suffix : str, optional - The suffix of the scope - """ - self.embedding_net_variables = get_embedding_net_variables(model_file, suffix = suffix) - - def prod_force_virial(self, atom_ener : tf.Tensor, natoms : tf.Tensor diff --git a/deepmd/descriptor/se_r.py b/deepmd/descriptor/se_r.py index 41abc4c677..e874a659b0 100644 --- a/deepmd/descriptor/se_r.py +++ b/deepmd/descriptor/se_r.py @@ -12,10 +12,12 @@ from deepmd.utils.graph import get_embedding_net_variables from deepmd.utils.sess import run_sess from .descriptor import Descriptor +from .se import DescrptSe + @Descriptor.register("se_e2_r") @Descriptor.register("se_r") -class DescrptSeR (Descriptor): +class DescrptSeR (DescrptSe): """DeepPot-SE constructed from radial information of atomic configurations. The embedding takes the distance between atoms as input. @@ -314,10 +316,7 @@ def build (self, sel = self.sel_r) self.descrpt_reshape = tf.reshape(self.descrpt, [-1, self.ndescrpt]) - self.descrpt_reshape = tf.identity(self.descrpt_reshape, name = 'o_rmat' + suffix) - self.descrpt_deriv = tf.identity(self.descrpt_deriv, name = 'o_rmat_deriv' + suffix) - self.rij = tf.identity(self.rij, name = 'o_rij' + suffix) - self.nlist = tf.identity(self.nlist, name = 'o_nlist' + suffix) + self._identity_tensors(suffix=suffix) # only used when tensorboard was set as true tf.summary.histogram('descrpt', self.descrpt) @@ -506,59 +505,3 @@ def _filter_r(self, result = tf.reduce_mean(xyz_scatter, axis = 1) * res_rescale return result - - def get_tensor_names(self, suffix : str = "") -> Tuple[str]: - """Get names of tensors. - - Parameters - ---------- - suffix : str - The suffix of the scope - - Returns - ------- - Tuple[str] - Names of tensors - """ - return (f'o_rmat{suffix}:0', f'o_rmat_deriv{suffix}:0', f'o_rij{suffix}:0', f'o_nlist{suffix}:0') - - def pass_tensors_from_frz_model(self, - descrpt_reshape : tf.Tensor, - descrpt_deriv : tf.Tensor, - rij : tf.Tensor, - nlist : tf.Tensor - ): - """ - Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def - - Parameters - ---------- - descrpt_reshape - The passed descrpt_reshape tensor - descrpt_deriv - The passed descrpt_deriv tensor - rij - The passed rij tensor - nlist - The passed nlist tensor - """ - self.rij = rij - self.nlist = nlist - self.descrpt_deriv = descrpt_deriv - self.descrpt_reshape = descrpt_reshape - - def init_variables(self, - model_file : str, - suffix : str = "", - ) -> None: - """ - Init the embedding net variables with the given dict - - Parameters - ---------- - model_file : str - The input frozen model file - suffix : str, optional - The suffix of the scope - """ - self.embedding_net_variables = get_embedding_net_variables(model_file, suffix = suffix) diff --git a/deepmd/descriptor/se_t.py b/deepmd/descriptor/se_t.py index e2bfb3449c..cd1bbb61f8 100644 --- a/deepmd/descriptor/se_t.py +++ b/deepmd/descriptor/se_t.py @@ -9,14 +9,14 @@ from deepmd.env import op_module from deepmd.env import default_tf_session_config from deepmd.utils.network import embedding_net, embedding_net_rand_seed_shift -from deepmd.utils.graph import get_embedding_net_variables from deepmd.utils.sess import run_sess from .descriptor import Descriptor +from .se import DescrptSe @Descriptor.register("se_e3") @Descriptor.register("se_at") @Descriptor.register("se_a_3be") -class DescrptSeT (Descriptor): +class DescrptSeT (DescrptSe): """DeepPot-SE constructed from all information (both angular and radial) of atomic configurations. @@ -313,10 +313,7 @@ def build (self, sel_r = self.sel_r) self.descrpt_reshape = tf.reshape(self.descrpt, [-1, self.ndescrpt]) - self.descrpt_reshape = tf.identity(self.descrpt_reshape, name = 'o_rmat' + suffix) - self.descrpt_deriv = tf.identity(self.descrpt_deriv, name = 'o_rmat_deriv' + suffix) - self.rij = tf.identity(self.rij, name = 'o_rij' + suffix) - self.nlist = tf.identity(self.nlist, name = 'o_nlist' + suffix) + self._identity_tensors(suffix=suffix) self.dout, self.qmat = self._pass_filter(self.descrpt_reshape, atype, @@ -525,59 +522,3 @@ def _filter(self, else: result += res_ij return result, None - - def get_tensor_names(self, suffix : str = "") -> Tuple[str]: - """Get names of tensors. - - Parameters - ---------- - suffix : str - The suffix of the scope - - Returns - ------- - Tuple[str] - Names of tensors - """ - return (f'o_rmat{suffix}:0', f'o_rmat_deriv{suffix}:0', f'o_rij{suffix}:0', f'o_nlist{suffix}:0') - - def pass_tensors_from_frz_model(self, - descrpt_reshape : tf.Tensor, - descrpt_deriv : tf.Tensor, - rij : tf.Tensor, - nlist : tf.Tensor - ): - """ - Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def - - Parameters - ---------- - descrpt_reshape - The passed descrpt_reshape tensor - descrpt_deriv - The passed descrpt_deriv tensor - rij - The passed rij tensor - nlist - The passed nlist tensor - """ - self.rij = rij - self.nlist = nlist - self.descrpt_deriv = descrpt_deriv - self.descrpt_reshape = descrpt_reshape - - def init_variables(self, - model_file : str, - suffix : str = "", - ) -> None: - """ - Init the embedding net variables with the given dict - - Parameters - ---------- - model_file : str - The input frozen model file - suffix : str, optional - The suffix of the scope - """ - self.embedding_net_variables = get_embedding_net_variables(model_file, suffix = suffix) From 6b77aa2c01adcdd9c9e744315b64d538e3861c1c Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Mon, 13 Sep 2021 21:34:13 -0400 Subject: [PATCH 3/5] clean other changes --- deepmd/descriptor/se_r.py | 1 - deepmd/descriptor/se_t.py | 1 + 2 files changed, 1 insertion(+), 1 deletion(-) diff --git a/deepmd/descriptor/se_r.py b/deepmd/descriptor/se_r.py index e874a659b0..4d4a39074b 100644 --- a/deepmd/descriptor/se_r.py +++ b/deepmd/descriptor/se_r.py @@ -9,7 +9,6 @@ from deepmd.env import op_module from deepmd.env import default_tf_session_config from deepmd.utils.network import embedding_net, embedding_net_rand_seed_shift -from deepmd.utils.graph import get_embedding_net_variables from deepmd.utils.sess import run_sess from .descriptor import Descriptor from .se import DescrptSe diff --git a/deepmd/descriptor/se_t.py b/deepmd/descriptor/se_t.py index cd1bbb61f8..f43fb6f40c 100644 --- a/deepmd/descriptor/se_t.py +++ b/deepmd/descriptor/se_t.py @@ -522,3 +522,4 @@ def _filter(self, else: result += res_ij return result, None + From 0661cc5be84762dda06d626529bb8ec56acbc5b6 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Mon, 13 Sep 2021 21:47:55 -0400 Subject: [PATCH 4/5] modify `embedding_net_pattern` to support se_e3 --- deepmd/utils/graph.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepmd/utils/graph.py b/deepmd/utils/graph.py index 53fd05cc67..ab5a95a1c3 100644 --- a/deepmd/utils/graph.py +++ b/deepmd/utils/graph.py @@ -129,7 +129,7 @@ def get_embedding_net_nodes_from_graph_def(graph_def: tf.GraphDef, suffix: str = The embedding net nodes within the given tf.GraphDef object """ embedding_net_nodes = {} - embedding_net_pattern = f"filter_type_\d+{suffix}/matrix_\d+_\d+|filter_type_\d+{suffix}/bias_\d+_\d+|filter_type_\d+{suffix}/idt_\d+_\d+|filter_type_all{suffix}/matrix_\d+_\d+|filter_type_all{suffix}/bias_\d+_\d+|filter_type_all{suffix}/idt_\d+_\d" + embedding_net_pattern = f"filter_type_\d+{suffix}/matrix_\d+_\d+|filter_type_\d+{suffix}/bias_\d+_\d+|filter_type_\d+{suffix}/idt_\d+_\d+|filter_type_all{suffix}/matrix_\d+_\d+|filter_type_all{suffix}/matrix_\d+_\d+_\d+|filter_type_all{suffix}/bias_\d+_\d+|filter_type_all{suffix}/bias_\d+_\d+_\d+|filter_type_all{suffix}/idt_\d+_\d" for node in graph_def.node: if re.fullmatch(embedding_net_pattern, node.name) != None: embedding_net_nodes[node.name] = node.attr["value"].tensor From b77ad4d2b1f8cd4b34421f3ee72ac1c618b73d29 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Tue, 14 Sep 2021 18:51:15 -0400 Subject: [PATCH 5/5] fix typo in `embedding_net_pattern` --- deepmd/utils/graph.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepmd/utils/graph.py b/deepmd/utils/graph.py index ab5a95a1c3..6766750e4b 100644 --- a/deepmd/utils/graph.py +++ b/deepmd/utils/graph.py @@ -129,7 +129,7 @@ def get_embedding_net_nodes_from_graph_def(graph_def: tf.GraphDef, suffix: str = The embedding net nodes within the given tf.GraphDef object """ embedding_net_nodes = {} - embedding_net_pattern = f"filter_type_\d+{suffix}/matrix_\d+_\d+|filter_type_\d+{suffix}/bias_\d+_\d+|filter_type_\d+{suffix}/idt_\d+_\d+|filter_type_all{suffix}/matrix_\d+_\d+|filter_type_all{suffix}/matrix_\d+_\d+_\d+|filter_type_all{suffix}/bias_\d+_\d+|filter_type_all{suffix}/bias_\d+_\d+_\d+|filter_type_all{suffix}/idt_\d+_\d" + embedding_net_pattern = f"filter_type_\d+{suffix}/matrix_\d+_\d+|filter_type_\d+{suffix}/bias_\d+_\d+|filter_type_\d+{suffix}/idt_\d+_\d+|filter_type_all{suffix}/matrix_\d+_\d+|filter_type_all{suffix}/matrix_\d+_\d+_\d+|filter_type_all{suffix}/bias_\d+_\d+|filter_type_all{suffix}/bias_\d+_\d+_\d+|filter_type_all{suffix}/idt_\d+_\d+" for node in graph_def.node: if re.fullmatch(embedding_net_pattern, node.name) != None: embedding_net_nodes[node.name] = node.attr["value"].tensor