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

* enable init-frz-model support for the original model within the dp train interface

* add init_variables method for ABC

* add doc for embedding_net method

Author: denghuilu

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__)

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

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)

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
@@ -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
     ----------
@@ -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)