From d5dd06ad296dbd5bbc16bce34a2b1c03fc5e0417 Mon Sep 17 00:00:00 2001 From: denghuilu Date: Tue, 23 Nov 2021 00:03:46 +0800 Subject: [PATCH 1/2] fix model compression bug when fparam and aparam are not zero --- deepmd/fit/ener.py | 33 +++++++++++++++++++++++++++++---- deepmd/train/trainer.py | 3 +++ 2 files changed, 32 insertions(+), 4 deletions(-) diff --git a/deepmd/fit/ener.py b/deepmd/fit/ener.py index 0afcf26de2..4915131638 100644 --- a/deepmd/fit/ener.py +++ b/deepmd/fit/ener.py @@ -9,7 +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.utils.graph import get_fitting_net_variables, load_graph_def, get_tensor_by_name_from_graph from deepmd.env import global_cvt_2_tf_float from deepmd.env import GLOBAL_TF_FLOAT_PRECISION @@ -150,6 +150,7 @@ def __init__ (self, self.aparam_inv_std = None self.fitting_net_variables = None + self.compress = False def get_numb_fparam(self) -> int: """ @@ -359,9 +360,9 @@ def build (self, The system energy """ bias_atom_e = self.bias_atom_e - if self.numb_fparam > 0 and ( self.fparam_avg is None or self.fparam_inv_std is None ): + if self.numb_fparam > 0 and ( self.fparam_avg is None or self.fparam_inv_std is None ) and not self.compress: raise RuntimeError('No data stat result. one should do data statisitic, before build') - if self.numb_aparam > 0 and ( self.aparam_avg is None or self.aparam_inv_std is None ): + if self.numb_aparam > 0 and ( self.aparam_avg is None or self.aparam_inv_std is None ) and not self.compress: raise RuntimeError('No data stat result. one should do data statisitic, before build') with tf.variable_scope('fitting_attr' + suffix, reuse = reuse) : @@ -494,4 +495,28 @@ def init_variables(self, model_file : str The input frozen model file """ - self.fitting_net_variables = get_fitting_net_variables(model_file) \ No newline at end of file + self.fitting_net_variables = get_fitting_net_variables(model_file) + + + def enable_compression(self, + model_file: str, + suffix: str = "" + ) -> None: + """ + Set the fitting net attributes from the frozen model_file when fparam or aparam is not zero + + Parameters + ---------- + model_file : str + The input frozen model file + suffix : str, optional + The suffix of the scope + """ + self.compress = True + graph, _ = load_graph_def(model_file) + if self.numb_fparam > 0: + self.fparam_avg = get_tensor_by_name_from_graph(graph, 'fitting_attr%s/t_fparam_avg' % suffix) + self.fparam_inv_std = get_tensor_by_name_from_graph(graph, 'fitting_attr%s/t_fparam_istd' % suffix) + if self.numb_aparam > 0: + self.aparam_avg = get_tensor_by_name_from_graph(graph, 'fitting_attr%s/t_aparam_avg' % suffix) + self.aparam_inv_std = get_tensor_by_name_from_graph(graph, 'fitting_attr%s/t_aparam_istd' % suffix) \ No newline at end of file diff --git a/deepmd/train/trainer.py b/deepmd/train/trainer.py index 16d1234112..e8b9dd1cef 100644 --- a/deepmd/train/trainer.py +++ b/deepmd/train/trainer.py @@ -284,6 +284,9 @@ def build (self, 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(self.model_param['compress']['model_file']) + # for fparam or aparam settings in 'ener' type fitting net + if self.fitting_type == 'ener': + self.fitting.enable_compression(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) From eef74979b1f7d6719ea16e17dc2d6c75d039aa5b Mon Sep 17 00:00:00 2001 From: denghuilu Date: Tue, 23 Nov 2021 00:10:04 +0800 Subject: [PATCH 2/2] Update ener.py --- deepmd/fit/ener.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/deepmd/fit/ener.py b/deepmd/fit/ener.py index 4915131638..1cc2aeb7f6 100644 --- a/deepmd/fit/ener.py +++ b/deepmd/fit/ener.py @@ -150,7 +150,6 @@ def __init__ (self, self.aparam_inv_std = None self.fitting_net_variables = None - self.compress = False def get_numb_fparam(self) -> int: """ @@ -360,9 +359,9 @@ def build (self, The system energy """ bias_atom_e = self.bias_atom_e - if self.numb_fparam > 0 and ( self.fparam_avg is None or self.fparam_inv_std is None ) and not self.compress: + if self.numb_fparam > 0 and ( self.fparam_avg is None or self.fparam_inv_std is None ): raise RuntimeError('No data stat result. one should do data statisitic, before build') - if self.numb_aparam > 0 and ( self.aparam_avg is None or self.aparam_inv_std is None ) and not self.compress: + if self.numb_aparam > 0 and ( self.aparam_avg is None or self.aparam_inv_std is None ): raise RuntimeError('No data stat result. one should do data statisitic, before build') with tf.variable_scope('fitting_attr' + suffix, reuse = reuse) : @@ -512,8 +511,8 @@ def enable_compression(self, suffix : str, optional The suffix of the scope """ - self.compress = True - graph, _ = load_graph_def(model_file) + if self.numb_fparam > 0 or self.numb_aparam > 0: + graph, _ = load_graph_def(model_file) if self.numb_fparam > 0: self.fparam_avg = get_tensor_by_name_from_graph(graph, 'fitting_attr%s/t_fparam_avg' % suffix) self.fparam_inv_std = get_tensor_by_name_from_graph(graph, 'fitting_attr%s/t_fparam_istd' % suffix)