Skip to content
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

Jump to bottom

enable model compression for dipole and polar type fitting net #1228

Merged
merged 2 commits into from
Oct 21, 2021
Merged

enable model compression for dipole and polar type fitting net #1228

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
b242c47
enable model compression for dipole and polar type fitting net
denghuilu Oct 20, 2021
db5c821
delete debug code
denghuilu Oct 21, 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
2 changes: 1 addition & 1 deletion deepmd/descriptor/se.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
----------
Expand Down
21 changes: 18 additions & 3 deletions deepmd/fit/dipole.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
"""
Expand Down Expand Up @@ -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])
Expand All @@ -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)
13 changes: 6 additions & 7 deletions deepmd/fit/ener.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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
self.fitting_net_variables = get_fitting_net_variables(model_file)
36 changes: 32 additions & 4 deletions deepmd/fit/polar.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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]:
"""
Expand Down Expand Up @@ -322,17 +324,17 @@ 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)
# bavg[0] = self.avgeig[0]
# 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])
Expand All @@ -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])
Expand All @@ -371,6 +373,19 @@ 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)


class GlobalPolarFittingSeA () :
"""
Expand Down Expand Up @@ -480,4 +495,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)

8 changes: 4 additions & 4 deletions deepmd/train/trainer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)