From b242c47584cd6fa94ea1e910183ac7530b6a3163 Mon Sep 17 00:00:00 2001 From: denghuilu Date: Wed, 20 Oct 2021 21:46:26 +0800 Subject: [PATCH 1/2] enable model compression for dipole and polar type fitting net --- deepmd/descriptor/se.py | 2 +- deepmd/fit/dipole.py | 21 ++++++++++++++++++--- deepmd/fit/ener.py | 13 ++++++------- deepmd/fit/polar.py | 37 +++++++++++++++++++++++++++++++++---- deepmd/train/trainer.py | 8 ++++---- 5 files changed, 62 insertions(+), 19 deletions(-) diff --git a/deepmd/descriptor/se.py b/deepmd/descriptor/se.py index 72bf908ecb..c7470568b6 100644 --- a/deepmd/descriptor/se.py +++ b/deepmd/descriptor/se.py @@ -96,7 +96,7 @@ def init_variables(self, suffix : str = "", ) -> None: """ - Init the embedding net variables with the given dict + Init the embedding net variables with the given frozen model Parameters ---------- diff --git a/deepmd/fit/dipole.py b/deepmd/fit/dipole.py index 81c26db9c9..6c115e3fb3 100644 --- a/deepmd/fit/dipole.py +++ b/deepmd/fit/dipole.py @@ -6,6 +6,7 @@ from deepmd.common import add_data_requirement, get_activation_func, get_precision, ACTIVATION_FN_DICT, PRECISION_DICT, docstring_parameter from deepmd.utils.argcheck import list_to_doc from deepmd.utils.network import one_layer, one_layer_rand_seed_shift +from deepmd.utils.graph import get_fitting_net_variables from deepmd.descriptor import DescrptSeA from deepmd.env import global_cvt_2_tf_float @@ -75,6 +76,7 @@ def __init__ (self, self.dim_rot_mat_1 = descrpt.get_dim_rot_mat_1() self.dim_rot_mat = self.dim_rot_mat_1 * 3 self.useBN = False + self.fitting_net_variables = None def get_sel_type(self) -> int: """ @@ -139,12 +141,12 @@ def build (self, layer = inputs_i for ii in range(0,len(self.n_neuron)) : if ii >= 1 and self.n_neuron[ii] == self.n_neuron[ii-1] : - layer+= one_layer(layer, self.n_neuron[ii], name='layer_'+str(ii)+'_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, use_timestep = self.resnet_dt, activation_fn = self.fitting_activation_fn, precision = self.fitting_precision, uniform_seed = self.uniform_seed) + layer+= one_layer(layer, self.n_neuron[ii], name='layer_'+str(ii)+'_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, use_timestep = self.resnet_dt, activation_fn = self.fitting_activation_fn, precision = self.fitting_precision, uniform_seed = self.uniform_seed, initial_variables = self.fitting_net_variables) else : - layer = one_layer(layer, self.n_neuron[ii], name='layer_'+str(ii)+'_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, activation_fn = self.fitting_activation_fn, precision = self.fitting_precision, uniform_seed = self.uniform_seed) + layer = one_layer(layer, self.n_neuron[ii], name='layer_'+str(ii)+'_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, activation_fn = self.fitting_activation_fn, precision = self.fitting_precision, uniform_seed = self.uniform_seed, initial_variables = self.fitting_net_variables) if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift # (nframes x natoms) x naxis - final_layer = one_layer(layer, self.dim_rot_mat_1, activation_fn = None, name='final_layer_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, precision = self.fitting_precision, uniform_seed = self.uniform_seed) + final_layer = one_layer(layer, self.dim_rot_mat_1, activation_fn = None, name='final_layer_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, precision = self.fitting_precision, uniform_seed = self.uniform_seed, initial_variables = self.fitting_net_variables) if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift # (nframes x natoms) x 1 * naxis final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0] * natoms[2+type_i], 1, self.dim_rot_mat_1]) @@ -163,3 +165,16 @@ def build (self, tf.summary.histogram('fitting_net_output', outs) return tf.cast(tf.reshape(outs, [-1]), GLOBAL_TF_FLOAT_PRECISION) # return tf.reshape(outs, [tf.shape(inputs)[0] * natoms[0] * 3 // 3]) + + def init_variables(self, + model_file: str + ) -> None: + """ + Init the fitting net variables with the given frozen model + + Parameters + ---------- + model_file : str + The input frozen model file + """ + self.fitting_net_variables = get_fitting_net_variables(model_file) \ No newline at end of file diff --git a/deepmd/fit/ener.py b/deepmd/fit/ener.py index 37900a70d8..0afcf26de2 100644 --- a/deepmd/fit/ener.py +++ b/deepmd/fit/ener.py @@ -9,6 +9,7 @@ from deepmd.descriptor import DescrptLocFrame from deepmd.descriptor import DescrptSeA from deepmd.utils.type_embed import embed_atom_type +from deepmd.utils.graph import get_fitting_net_variables from deepmd.env import global_cvt_2_tf_float from deepmd.env import GLOBAL_TF_FLOAT_PRECISION @@ -148,7 +149,6 @@ def __init__ (self, self.aparam_std = None self.aparam_inv_std = None - self.compress = False self.fitting_net_variables = None def get_numb_fparam(self) -> int: @@ -484,15 +484,14 @@ def build (self, def init_variables(self, - fitting_net_variables: dict + model_file: str ) -> None: """ - Init the fitting net variables with the given dict + Init the fitting net variables with the given frozen model Parameters ---------- - fitting_net_variables - The input dict which stores the fitting net variables + model_file : str + The input frozen model file """ - self.compress = True - self.fitting_net_variables = fitting_net_variables \ No newline at end of file + self.fitting_net_variables = get_fitting_net_variables(model_file) \ No newline at end of file diff --git a/deepmd/fit/polar.py b/deepmd/fit/polar.py index c82e08f214..49404a1264 100644 --- a/deepmd/fit/polar.py +++ b/deepmd/fit/polar.py @@ -6,6 +6,7 @@ from deepmd.common import add_data_requirement, get_activation_func, get_precision, ACTIVATION_FN_DICT, PRECISION_DICT, docstring_parameter from deepmd.utils.argcheck import list_to_doc from deepmd.utils.network import one_layer, one_layer_rand_seed_shift +from deepmd.utils.graph import get_fitting_net_variables from deepmd.descriptor import DescrptLocFrame from deepmd.descriptor import DescrptSeA @@ -192,6 +193,7 @@ def __init__ (self, self.dim_rot_mat_1 = descrpt.get_dim_rot_mat_1() self.dim_rot_mat = self.dim_rot_mat_1 * 3 self.useBN = False + self.fitting_net_variables = None def get_sel_type(self) -> List[int]: """ @@ -322,9 +324,9 @@ def build (self, layer = inputs_i for ii in range(0,len(self.n_neuron)) : if ii >= 1 and self.n_neuron[ii] == self.n_neuron[ii-1] : - layer+= one_layer(layer, self.n_neuron[ii], name='layer_'+str(ii)+'_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, use_timestep = self.resnet_dt, activation_fn = self.fitting_activation_fn, precision = self.fitting_precision, uniform_seed = self.uniform_seed) + layer+= one_layer(layer, self.n_neuron[ii], name='layer_'+str(ii)+'_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, use_timestep = self.resnet_dt, activation_fn = self.fitting_activation_fn, precision = self.fitting_precision, uniform_seed = self.uniform_seed, initial_variables = self.fitting_net_variables) else : - layer = one_layer(layer, self.n_neuron[ii], name='layer_'+str(ii)+'_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, activation_fn = self.fitting_activation_fn, precision = self.fitting_precision, uniform_seed = self.uniform_seed) + layer = one_layer(layer, self.n_neuron[ii], name='layer_'+str(ii)+'_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, activation_fn = self.fitting_activation_fn, precision = self.fitting_precision, uniform_seed = self.uniform_seed, initial_variables = self.fitting_net_variables) if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift if self.fit_diag : bavg = np.zeros(self.dim_rot_mat_1) @@ -332,7 +334,7 @@ def build (self, # bavg[1] = self.avgeig[1] # bavg[2] = self.avgeig[2] # (nframes x natoms) x naxis - final_layer = one_layer(layer, self.dim_rot_mat_1, activation_fn = None, name='final_layer_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, bavg = bavg, precision = self.fitting_precision, uniform_seed = self.uniform_seed) + final_layer = one_layer(layer, self.dim_rot_mat_1, activation_fn = None, name='final_layer_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, bavg = bavg, precision = self.fitting_precision, uniform_seed = self.uniform_seed, initial_variables = self.fitting_net_variables) if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift # (nframes x natoms) x naxis final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0] * natoms[2+type_i], self.dim_rot_mat_1]) @@ -344,7 +346,7 @@ def build (self, # bavg[1*self.dim_rot_mat_1+1] = self.avgeig[1] # bavg[2*self.dim_rot_mat_1+2] = self.avgeig[2] # (nframes x natoms) x (naxis x naxis) - final_layer = one_layer(layer, self.dim_rot_mat_1*self.dim_rot_mat_1, activation_fn = None, name='final_layer_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, bavg = bavg, precision = self.fitting_precision, uniform_seed = self.uniform_seed) + final_layer = one_layer(layer, self.dim_rot_mat_1*self.dim_rot_mat_1, activation_fn = None, name='final_layer_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, bavg = bavg, precision = self.fitting_precision, uniform_seed = self.uniform_seed, initial_variables = self.fitting_net_variables) if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift # (nframes x natoms) x naxis x naxis final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0] * natoms[2+type_i], self.dim_rot_mat_1, self.dim_rot_mat_1]) @@ -371,6 +373,20 @@ def build (self, tf.summary.histogram('fitting_net_output', outs) return tf.cast(tf.reshape(outs, [-1]), GLOBAL_TF_FLOAT_PRECISION) + def init_variables(self, + model_file: str + ) -> None: + """ + Init the fitting net variables with the given frozen model + + Parameters + ---------- + model_file : str + The input frozen model file + """ + self.fitting_net_variables = get_fitting_net_variables(model_file) + print(self.fitting_net_variables.keys()) + class GlobalPolarFittingSeA () : """ @@ -480,4 +496,17 @@ def build (self, outs = tf.reduce_sum(outs, axis = 1) tf.summary.histogram('fitting_net_output', outs) return tf.reshape(outs, [-1]) + + def init_variables(self, + model_file: str + ) -> None: + """ + Init the fitting net variables with the given frozen model + + Parameters + ---------- + model_file : str + The input frozen model file + """ + self.polar_fitting.init_variables(model_file) diff --git a/deepmd/train/trainer.py b/deepmd/train/trainer.py index a8d85d86e6..16d1234112 100644 --- a/deepmd/train/trainer.py +++ b/deepmd/train/trainer.py @@ -20,7 +20,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_embedding_net_variables, get_fitting_net_variables +from deepmd.utils.graph import get_tensor_by_name from tensorflow.python.client import timeline from deepmd.env import op_module @@ -283,7 +283,7 @@ def build (self, # architecture to call neighbor stat else : self.descrpt.enable_compression(self.model_param['compress']["min_nbor_dist"], self.model_param['compress']['model_file'], self.model_param['compress']['table_config'][0], self.model_param['compress']['table_config'][1], self.model_param['compress']['table_config'][2], self.model_param['compress']['table_config'][3]) - self.fitting.init_variables(get_fitting_net_variables(self.model_param['compress']['model_file'])) + self.fitting.init_variables(self.model_param['compress']['model_file']) if self.is_compress or self.model_type == 'compressed_model': tf.constant("compressed_model", name = 'model_type', dtype = tf.string) @@ -661,11 +661,11 @@ def _init_from_frz_model(self): # initialize fitting net with the given compressed frozen model if self.model_type == 'original_model': self.descrpt.init_variables(self.run_opt.init_frz_model) - self.fitting.init_variables(get_fitting_net_variables(self.run_opt.init_frz_model)) + self.fitting.init_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)) + self.fitting.init_variables(self.frz_model) tf.constant("compressed_model", name = 'model_type', dtype = tf.string) else: raise RuntimeError("Unknown model type %s" % self.model_type) From db5c8216d6cc5631e616db46144b50270547271e Mon Sep 17 00:00:00 2001 From: denghuilu Date: Thu, 21 Oct 2021 10:41:25 +0800 Subject: [PATCH 2/2] delete debug code --- deepmd/fit/polar.py | 1 - 1 file changed, 1 deletion(-) diff --git a/deepmd/fit/polar.py b/deepmd/fit/polar.py index 49404a1264..65b1ff6aef 100644 --- a/deepmd/fit/polar.py +++ b/deepmd/fit/polar.py @@ -385,7 +385,6 @@ def init_variables(self, The input frozen model file """ self.fitting_net_variables = get_fitting_net_variables(model_file) - print(self.fitting_net_variables.keys()) class GlobalPolarFittingSeA () :