Skip to content

Enable init-frz-model support for the original model #1102

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
amcadmus merged 3 commits into from
Sep 6, 2021
Merged

Enable init-frz-model support for the original model #1102

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
f0fd030
enable init-frz-model support for the original model within the dp tr…
denghuilu Sep 6, 2021
cd1c7ee
add init_variables method for ABC
denghuilu Sep 6, 2021
0750177
add doc for embedding_net method
denghuilu Sep 6, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 20 additions & 0 deletions deepmd/descriptor/descriptor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -283,3 +283,23 @@ def get_feed_dict(self,
# TODO: currently only SeA has this method, but I think the method can be
# moved here as it doesn't contain anything related to a specific descriptor
raise NotImplementedError

def init_variables(self,
embedding_net_variables: dict
) -> None:
"""
Init the embedding net variables with the given dict

Parameters
----------
embedding_net_variables
The input dict which stores the embedding net variables

Notes
-----
This method is called by others when the descriptor supported initialization from the given variables.
"""
# TODO: currently only SeA has this method, but I think the method can be
# moved here as it doesn't contain anything related to a specific descriptor
raise NotImplementedError(
"Descriptor %s doesn't support initialization from the given variables!" % type(self).__name__)
19 changes: 18 additions & 1 deletion deepmd/descriptor/se_a.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -157,6 +157,7 @@ def __init__ (self,
self.dstd = None
self.davg = None
self.compress = False
self.embedding_net_variables = None
self.place_holders = {}
nei_type = np.array([])
for ii in range(self.ntypes):
Expand Down Expand Up @@ -521,6 +522,21 @@ def get_feed_dict(self,
}
return feed_dict


def init_variables(self,
embedding_net_variables: dict
) -> None:
"""
Init the embedding net variables with the given dict

Parameters
----------
embedding_net_variables
The input dict which stores the embedding net variables
"""
self.embedding_net_variables = embedding_net_variables


def prod_force_virial(self,
atom_ener : tf.Tensor,
natoms : tf.Tensor
Expand Down Expand Up @@ -766,7 +782,8 @@ def _filter_lower(
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
else:
# we can safely return the final xyz_scatter filled with zero directly
Expand Down
16 changes: 9 additions & 7 deletions deepmd/train/trainer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@
from deepmd.utils.neighbor_stat import NeighborStat
from deepmd.utils.sess import run_sess
from deepmd.utils.type_embed import TypeEmbedNet
from deepmd.utils.graph import get_tensor_by_name, get_fitting_net_variables
from deepmd.utils.graph import get_tensor_by_name, get_embedding_net_variables, get_fitting_net_variables

from tensorflow.python.client import timeline
from deepmd.env import op_module
Expand Down Expand Up @@ -278,7 +278,6 @@ def _init_param(self, jdata):
# if init the graph with the frozen model
self.frz_model = None
self.model_type = None
self.init_from_frz_model = False


def build (self,
Expand Down Expand Up @@ -694,14 +693,17 @@ def _init_from_frz_model(self):
"which is not supported by the 'dp train init-frz-model' interface. " % self.run_opt.init_frz_model
) from e

if self.fitting_type != 'ener':
raise RuntimeError("The 'dp train init-frz-model' command only supports the 'ener' type fitting net currently!")
# self.frz_model will control the self.model to import the descriptor from the given frozen model instead of building from scratch...
# initialize fitting net with the given compressed frozen model
if self.model_type == 'compressed_model' and self.fitting_type == 'ener':
self.init_from_frz_model = True
if self.model_type == 'original_model':
self.descrpt.init_variables(get_embedding_net_variables(self.run_opt.init_frz_model))
self.fitting.init_variables(get_fitting_net_variables(self.run_opt.init_frz_model))
tf.constant("original_model", name = 'model_type', dtype = tf.string)
elif self.model_type == 'compressed_model':
self.frz_model = self.run_opt.init_frz_model
self.fitting.init_variables(get_fitting_net_variables(self.frz_model))
tf.constant("compressed_model", name = 'model_type', dtype = tf.string)
elif self.fitting_type != 'ener':
raise RuntimeError("The 'dp train init-frz-model' command only supports the 'ener' type fitting net currently!")
else:
raise RuntimeError("The 'dp train init-frz-model' command only supports the compressed model currently!")
raise RuntimeError("Unknown model type %s" % self.model_type)
48 changes: 32 additions & 16 deletions deepmd/utils/network.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -92,7 +92,8 @@ def embedding_net(xx,
bavg = 0.0,
seed = None,
trainable = True,
uniform_seed = False):
uniform_seed = False,
initial_variables = None):
r"""The embedding network.

The embedding network function :math:`\mathcal{N}` is constructed by is the
Expand Down Expand Up @@ -141,6 +142,11 @@ def embedding_net(xx,
Random seed for initializing network parameters
trainable: boolean
If the network is trainable
uniform_seed : boolean
Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed
initial_variables : dict
The input dict which stores the embedding net variables


References
----------
Expand All @@ -152,37 +158,47 @@ def embedding_net(xx,
outputs_size = [input_shape[1]] + network_size

for ii in range(1, len(outputs_size)):
w = tf.get_variable('matrix_'+str(ii)+name_suffix,
w_initializer = tf.random_normal_initializer(
stddev=stddev/np.sqrt(outputs_size[ii]+outputs_size[ii-1]),
seed = seed if (seed is None or uniform_seed) else seed + ii*3+0
)
b_initializer = tf.random_normal_initializer(
stddev=stddev,
mean = bavg,
seed = seed if (seed is None or uniform_seed) else seed + 3*ii+1
)
if initial_variables is not None:
scope = tf.get_variable_scope().name
w_initializer = tf.constant_initializer(initial_variables[scope+'/matrix_'+str(ii)+name_suffix])
b_initializer = tf.constant_initializer(initial_variables[scope+'/bias_'+str(ii)+name_suffix])
w = tf.get_variable('matrix_'+str(ii)+name_suffix,
[outputs_size[ii - 1], outputs_size[ii]],
precision,
tf.random_normal_initializer(
stddev=stddev/np.sqrt(outputs_size[ii]+outputs_size[ii-1]),
seed = seed if (seed is None or uniform_seed) else seed + ii*3+0
),
w_initializer,
trainable = trainable)
variable_summaries(w, 'matrix_'+str(ii)+name_suffix)

b = tf.get_variable('bias_'+str(ii)+name_suffix,
[1, outputs_size[ii]],
precision,
tf.random_normal_initializer(
stddev=stddev,
mean = bavg,
seed = seed if (seed is None or uniform_seed) else seed + 3*ii+1
),
b_initializer,
trainable = trainable)
variable_summaries(b, 'bias_'+str(ii)+name_suffix)

hidden = tf.reshape(activation_fn(tf.matmul(xx, w) + b), [-1, outputs_size[ii]])
if resnet_dt :
idt_initializer = tf.random_normal_initializer(
stddev=0.001,
mean = 1.0,
seed = seed if (seed is None or uniform_seed) else seed + 3*ii+2
)
if initial_variables is not None:
scope = tf.get_variable_scope().name
idt_initializer = tf.constant_initializer(initial_variables[scope+'/idt_'+str(ii)+name_suffix])
idt = tf.get_variable('idt_'+str(ii)+name_suffix,
[1, outputs_size[ii]],
precision,
tf.random_normal_initializer(
stddev=0.001,
mean = 1.0,
seed = seed if (seed is None or uniform_seed) else seed + 3*ii+2
),
idt_initializer,
trainable = trainable)
variable_summaries(idt, 'idt_'+str(ii)+name_suffix)

Expand Down