Skip to content

Commit

Permalink
Add the missing initializations for extra embedding variables (#3005)
Browse files Browse the repository at this point in the history
  • Loading branch information
@nahso
nahso authored Nov 24, 2023
1 parent 4d1c5fd commit f321046
Show file tree
Hide file tree
Showing 5 changed files with 331 additions and 46 deletions.
41 changes: 28 additions & 13 deletions deepmd/descriptor/se_a.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,8 @@
GraphWithoutTensorError,
)
from deepmd.utils.graph import (
get_extra_embedding_net_suffix,
get_extra_embedding_net_variables_from_graph_def,
get_pattern_nodes_from_graph_def,
get_tensor_by_name_from_graph,
)
Expand Down Expand Up @@ -204,7 +206,7 @@ def __init__(
self.type_one_side = type_one_side
self.spin = spin
self.stripped_type_embedding = stripped_type_embedding
self.extra_embeeding_net_variables = None
self.extra_embedding_net_variables = None
self.layer_size = len(neuron)

# extend sel_a for spin system
Expand Down Expand Up @@ -470,11 +472,13 @@ def enable_compression(
)

if self.stripped_type_embedding:
one_side_suffix = get_extra_embedding_net_suffix(type_one_side=True)
two_side_suffix = get_extra_embedding_net_suffix(type_one_side=False)
ret_two_side = get_pattern_nodes_from_graph_def(
graph_def, f"filter_type_all{suffix}/.+_two_side_ebd"
graph_def, f"filter_type_all{suffix}/.+{two_side_suffix}"
)
ret_one_side = get_pattern_nodes_from_graph_def(
graph_def, f"filter_type_all{suffix}/.+_one_side_ebd"
graph_def, f"filter_type_all{suffix}/.+{one_side_suffix}"
)
if len(ret_two_side) == 0 and len(ret_one_side) == 0:
raise RuntimeError(
Expand All @@ -487,19 +491,19 @@ def enable_compression(
elif len(ret_two_side) != 0:
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, graph_def, two_side_suffix, "matrix", suffix
)
self.bias = get_extra_side_embedding_net_variable(
self, graph_def, "two_side", "bias", suffix
self, graph_def, two_side_suffix, "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, graph_def, one_side_suffix, "matrix", suffix
)
self.bias = get_extra_side_embedding_net_variable(
self, graph_def, "one_side", "bias", suffix
self, graph_def, one_side_suffix, "bias", suffix
)
self.extra_embedding = make_data(self, self.final_type_embedding)

Expand Down Expand Up @@ -961,20 +965,21 @@ def _filter_lower(

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,
name_suffix=get_extra_embedding_net_suffix(
self.type_one_side
),
stddev=stddev,
bavg=bavg,
seed=self.seed,
trainable=trainable,
uniform_seed=self.uniform_seed,
initial_variables=self.extra_embeeding_net_variables,
initial_variables=self.extra_embedding_net_variables,
mixed_prec=self.mixed_prec,
)
net_output = tf.nn.embedding_lookup(
Expand Down Expand Up @@ -1004,20 +1009,21 @@ def _filter_lower(
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,
name_suffix=get_extra_embedding_net_suffix(
self.type_one_side
),
stddev=stddev,
bavg=bavg,
seed=self.seed,
trainable=trainable,
uniform_seed=self.uniform_seed,
initial_variables=self.extra_embeeding_net_variables,
initial_variables=self.extra_embedding_net_variables,
mixed_prec=self.mixed_prec,
)
net_output = tf.nn.embedding_lookup(net_output, idx)
Expand Down Expand Up @@ -1327,6 +1333,15 @@ def init_variables(
self.dstd = new_dstd
if self.original_sel is None:
self.original_sel = sel
if self.stripped_type_embedding:
self.extra_embedding_net_variables = (
get_extra_embedding_net_variables_from_graph_def(
graph_def,
suffix,
get_extra_embedding_net_suffix(self.type_one_side),
self.layer_size,
)
)

@property
def explicit_ntypes(self) -> bool:
Expand Down
49 changes: 19 additions & 30 deletions deepmd/descriptor/se_atten.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -42,9 +42,10 @@
)
from deepmd.utils.graph import (
get_attention_layer_variables_from_graph_def,
get_extra_embedding_net_suffix,
get_extra_embedding_net_variables_from_graph_def,
get_pattern_nodes_from_graph_def,
get_tensor_by_name_from_graph,
get_tensor_by_type,
)
from deepmd.utils.network import (
embedding_net,
Expand Down Expand Up @@ -391,11 +392,12 @@ def enable_compression(
raise RuntimeError("can not compress model when attention layer is not 0.")

ret = get_pattern_nodes_from_graph_def(
graph_def, f"filter_type_all{suffix}/.+_two_side_ebd"
graph_def,
f"filter_type_all{suffix}/.+{get_extra_embedding_net_suffix(type_one_side=False)}",
)
if len(ret) == 0:
raise RuntimeError(
"can not find variables of embedding net `*_two_side_ebd` from graph_def, maybe it is not a compressible model."
f"can not find variables of embedding net `*{get_extra_embedding_net_suffix(type_one_side=False)}` from graph_def, maybe it is not a compressible model."
)

self.compress = True
Expand All @@ -420,11 +422,12 @@ def enable_compression(
)

self.final_type_embedding = get_two_side_type_embedding(self, graph)
type_side_suffix = get_extra_embedding_net_suffix(type_one_side=False)
self.matrix = get_extra_side_embedding_net_variable(
self, graph_def, "two_side", "matrix", suffix
self, graph_def, type_side_suffix, "matrix", suffix
)
self.bias = get_extra_side_embedding_net_variable(
self, graph_def, "two_side", "bias", suffix
self, graph_def, type_side_suffix, "bias", suffix
)
self.two_embd = make_data(self, self.final_type_embedding)

Expand Down Expand Up @@ -1125,14 +1128,15 @@ def _filter_lower(
two_side_type_embedding,
[-1, two_side_type_embedding.shape[-1]],
)
two_side_type_embedding_suffix = "_two_side_ebd"
embedding_of_two_side_type_embedding = embedding_net(
two_side_type_embedding,
self.filter_neuron,
self.filter_precision,
activation_fn=activation_fn,
resnet_dt=self.filter_resnet_dt,
name_suffix=two_side_type_embedding_suffix,
name_suffix=get_extra_embedding_net_suffix(
type_one_side=False
),
stddev=stddev,
bavg=bavg,
seed=self.seed,
Expand Down Expand Up @@ -1292,18 +1296,6 @@ def init_variables(
"""
super().init_variables(graph=graph, graph_def=graph_def, suffix=suffix)

if self.stripped_type_embedding:
self.two_side_embeeding_net_variables = {}
for i in range(1, self.layer_size + 1):
matrix_pattern = f"filter_type_all{suffix}/matrix_{i}_two_side_ebd"
self.two_side_embeeding_net_variables[
matrix_pattern
] = self._get_two_embed_variables(graph_def, matrix_pattern)
bias_pattern = f"filter_type_all{suffix}/bias_{i}_two_side_ebd"
self.two_side_embeeding_net_variables[
bias_pattern
] = self._get_two_embed_variables(graph_def, bias_pattern)

self.attention_layer_variables = get_attention_layer_variables_from_graph_def(
graph_def, suffix=suffix
)
Expand All @@ -1322,18 +1314,15 @@ def init_variables(
f"attention_layer_{i}{suffix}/layer_normalization_{i}/gamma"
]

def _get_two_embed_variables(self, graph_def, pattern: str):
node = get_pattern_nodes_from_graph_def(graph_def, pattern)[pattern]
dtype = tf.as_dtype(node.dtype).as_numpy_dtype
tensor_shape = tf.TensorShape(node.tensor_shape).as_list()
if (len(tensor_shape) != 1) or (tensor_shape[0] != 1):
tensor_value = np.frombuffer(
node.tensor_content,
dtype=tf.as_dtype(node.dtype).as_numpy_dtype,
if self.stripped_type_embedding:
self.two_side_embeeding_net_variables = (
get_extra_embedding_net_variables_from_graph_def(
graph_def,
suffix,
get_extra_embedding_net_suffix(type_one_side=False),
self.layer_size,
)
)
else:
tensor_value = get_tensor_by_type(node, dtype)
return np.reshape(tensor_value, tensor_shape)

def build_type_exclude_mask(
self,
Expand Down
6 changes: 3 additions & 3 deletions deepmd/utils/compress.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -43,15 +43,15 @@ def get_two_side_type_embedding(self, graph):


def get_extra_side_embedding_net_variable(
self, graph_def, type_side, varialbe_name, suffix
self, graph_def, type_side_suffix, 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",
f"filter_type_all{suffix}/{varialbe_name}_{i}{type_side_suffix}",
)
node = target[f"filter_type_all{suffix}/{varialbe_name}_{i}_{type_side}_ebd"]
node = target[f"filter_type_all{suffix}/{varialbe_name}_{i}{type_side_suffix}"]
ret["layer_" + str(i)] = node
return ret

Expand Down
85 changes: 85 additions & 0 deletions deepmd/utils/graph.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -237,6 +237,91 @@ def get_embedding_net_variables_from_graph_def(
return embedding_net_variables


def get_extra_embedding_net_suffix(type_one_side: bool):
"""Get the extra embedding net suffix according to the value of type_one_side.
Parameters
----------
type_one_side
The value of type_one_side
Returns
-------
str
The extra embedding net suffix
"""
if type_one_side:
extra_suffix = "_one_side_ebd"
else:
extra_suffix = "_two_side_ebd"
return extra_suffix


def get_variables_from_graph_def_as_numpy_array(graph_def: tf.GraphDef, pattern: str):
"""Get variables from the given tf.GraphDef object, with numpy array returns.
Parameters
----------
graph_def
The input tf.GraphDef object
pattern : str
The name of variable
Returns
-------
np.ndarray
The numpy array of the variable
"""
node = get_pattern_nodes_from_graph_def(graph_def, pattern)[pattern]
dtype = tf.as_dtype(node.dtype).as_numpy_dtype
tensor_shape = tf.TensorShape(node.tensor_shape).as_list()
if (len(tensor_shape) != 1) or (tensor_shape[0] != 1):
tensor_value = np.frombuffer(
node.tensor_content,
dtype=tf.as_dtype(node.dtype).as_numpy_dtype,
)
else:
tensor_value = get_tensor_by_type(node, dtype)
return np.reshape(tensor_value, tensor_shape)


def get_extra_embedding_net_variables_from_graph_def(
graph_def: tf.GraphDef, suffix: str, extra_suffix: str, layer_size: int
):
"""Get extra embedding net variables from the given tf.GraphDef object.
The "extra embedding net" means the embedding net with only type embeddings input,
which occurs in "se_atten_v2" and "se_a_ebd_v2" descriptor.
Parameters
----------
graph_def
The input tf.GraphDef object
suffix : str
The "common" suffix in the descriptor
extra_suffix : str
This value depends on the value of "type_one_side".
It should always be "_one_side_ebd" or "_two_side_ebd"
layer_size : int
The layer size of the embedding net
Returns
-------
Dict
The extra embedding net variables within the given tf.GraphDef object
"""
extra_embedding_net_variables = {}
for i in range(1, layer_size + 1):
matrix_pattern = f"filter_type_all{suffix}/matrix_{i}{extra_suffix}"
extra_embedding_net_variables[
matrix_pattern
] = get_variables_from_graph_def_as_numpy_array(graph_def, matrix_pattern)
bias_pattern = f"filter_type_all{suffix}/bias_{i}{extra_suffix}"
extra_embedding_net_variables[
bias_pattern
] = get_variables_from_graph_def_as_numpy_array(graph_def, bias_pattern)
return extra_embedding_net_variables


def get_embedding_net_variables(model_file: str, suffix: str = "") -> Dict:
"""Get the embedding net variables with the given frozen model(model_file).
Expand Down
Loading

0 comments on commit f321046

Please sign in to comment.