From f0fd030b40ee43bbb252051a8cc22ffcfd81764c Mon Sep 17 00:00:00 2001 From: denghuilu Date: Mon, 6 Sep 2021 11:03:53 +0800 Subject: [PATCH 1/3] enable init-frz-model support for the original model within the dp train interface --- deepmd/descriptor/se_a.py | 19 ++++++++++++++++- deepmd/train/trainer.py | 16 ++++++++------- deepmd/utils/network.py | 43 ++++++++++++++++++++++++--------------- 3 files changed, 54 insertions(+), 24 deletions(-) diff --git a/deepmd/descriptor/se_a.py b/deepmd/descriptor/se_a.py index da600848fb..47e28522a9 100644 --- a/deepmd/descriptor/se_a.py +++ b/deepmd/descriptor/se_a.py @@ -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): @@ -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 @@ -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 diff --git a/deepmd/train/trainer.py b/deepmd/train/trainer.py index 0fa08e77a2..f81f582e24 100644 --- a/deepmd/train/trainer.py +++ b/deepmd/train/trainer.py @@ -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 @@ -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, @@ -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) diff --git a/deepmd/utils/network.py b/deepmd/utils/network.py index d42f3a5c48..346ee867f2 100644 --- a/deepmd/utils/network.py +++ b/deepmd/utils/network.py @@ -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 @@ -152,37 +153,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) From cd1c7eecdb86fa27af92de54ed216294912eb3eb Mon Sep 17 00:00:00 2001 From: denghuilu Date: Mon, 6 Sep 2021 11:22:21 +0800 Subject: [PATCH 2/3] add init_variables method for ABC --- deepmd/descriptor/descriptor.py | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/deepmd/descriptor/descriptor.py b/deepmd/descriptor/descriptor.py index 4516aa146c..f241f6ee15 100644 --- a/deepmd/descriptor/descriptor.py +++ b/deepmd/descriptor/descriptor.py @@ -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__) From 07501775d735aa8108d3edee470f60aadbe5fd9f Mon Sep 17 00:00:00 2001 From: denghuilu Date: Mon, 6 Sep 2021 14:01:13 +0800 Subject: [PATCH 3/3] add doc for embedding_net method --- deepmd/utils/network.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/deepmd/utils/network.py b/deepmd/utils/network.py index 346ee867f2..5c78031167 100644 --- a/deepmd/utils/network.py +++ b/deepmd/utils/network.py @@ -142,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 ----------