def enable_compression(
self,
min_nbor_dist: float,
model_file: str = 'frozon_model.pb',
table_extrapolate: float = 5,
table_stride_1: float = 0.01,
table_stride_2: float = 0.1,
check_frequency: int = -1,
suffix: str = "",
) -> None:
"""
Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
Parameters
----------
min_nbor_dist
The nearest distance between atoms
model_file
The original frozen model, which will be compressed by the program
table_extrapolate
The scale of model extrapolation
table_stride_1
The uniform stride of the first table
table_stride_2
The uniform stride of the second table
check_frequency
The overflow check frequency
suffix : str, optional
The suffix of the scope
"""
assert (
not self.filter_resnet_dt
), "Model compression error: descriptor resnet_dt must be false!"
self.compress = True
self.table = DPTabulate(model_file,
self.type_one_side,
self.exclude_types,
self.compress_activation_fn,
suffix=suffix)
self.table_config = [
table_extrapolate, table_stride_1, table_stride_2, check_frequency
]
self.lower, self.upper \
= self.table.build(min_nbor_dist,
table_extrapolate,
table_stride_1,
table_stride_2)
graph, _ = load_graph_def(model_file)
self.davg = get_tensor_by_name_from_graph(
graph, 'descrpt_attr%s/t_avg' % suffix)
self.dstd = get_tensor_by_name_from_graph(
graph, 'descrpt_attr%s/t_std' % suffix)
def init_variables(self,
graph: tf.Graph,
graph_def: tf.GraphDef,
suffix : str = "",
) -> None:
"""
Init the embedding net variables with the given dict
Parameters
----------
graph : tf.Graph
The input frozen model graph
graph_def : tf.GraphDef
The input frozen model graph_def
suffix : str, optional
The suffix of the scope
"""
self.embedding_net_variables = get_embedding_net_variables_from_graph_def(graph_def, suffix = suffix)
self.davg = get_tensor_by_name_from_graph(graph, 'descrpt_attr%s/t_avg' % suffix)
self.dstd = get_tensor_by_name_from_graph(graph, 'descrpt_attr%s/t_std' % suffix)
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
"""
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)
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)
def _init_from_frz_model(self):
try:
graph, graph_def = load_graph_def(self.run_opt.init_frz_model)
except FileNotFoundError as e:
# throw runtime error if there's no frozen model
raise RuntimeError(
"The input frozen model %s (%s) does not exist! Please check the path of the frozen model. "
% (self.run_opt.init_frz_model,
os.path.abspath(self.run_opt.init_frz_model))) from e
# get the model type from the frozen model(self.run_opt.init_frz_model)
try:
t_model_type = get_tensor_by_name_from_graph(graph, 'model_type')
except GraphWithoutTensorError as e:
# throw runtime error if the frozen_model has no model type information...
raise RuntimeError(
"The input frozen model: %s has no 'model_type' information, "
"which is not supported by the 'dp train init-frz-model' interface. "
% self.run_opt.init_frz_model) from e
else:
self.model_type = bytes.decode(t_model_type)
self.model.init_variables(graph, graph_def, model_type=self.model_type)
def __init__(
self,
model_file: str,
type_one_side: bool = False,
exclude_types: List[List[int]] = [],
activation_fn: Callable[[tf.Tensor], tf.Tensor] = tf.nn.tanh,
suffix: str = "",
) -> None:
"""
Constructor
"""
self.model_file = model_file
self.type_one_side = type_one_side
self.exclude_types = exclude_types
self.suffix = suffix
if self.type_one_side and len(self.exclude_types) != 0:
raise RuntimeError(
'"type_one_side" is not compatible with "exclude_types"')
# functype
if activation_fn == ACTIVATION_FN_DICT["tanh"]:
self.functype = 1
elif activation_fn == ACTIVATION_FN_DICT["gelu"]:
self.functype = 2
else:
raise RuntimeError("Unknown actication function type!")
self.activation_fn = activation_fn
self.graph, self.graph_def = load_graph_def(self.model_file)
self.sess = tf.Session(graph=self.graph)
self.sub_graph, self.sub_graph_def = self._load_sub_graph()
self.sub_sess = tf.Session(graph=self.sub_graph)
try:
self.sel_a = self.graph.get_operation_by_name(
'ProdEnvMatA').get_attr('sel_a')
self.descrpt = self.graph.get_operation_by_name('ProdEnvMatA')
except Exception:
self.sel_a = self.graph.get_operation_by_name(
'DescrptSeA').get_attr('sel_a')
self.descrpt = self.graph.get_operation_by_name('DescrptSeA')
self.davg = get_tensor_by_name_from_graph(
self.graph, f'descrpt_attr{self.suffix}/t_avg')
self.dstd = get_tensor_by_name_from_graph(
self.graph, f'descrpt_attr{self.suffix}/t_std')
self.ntypes = get_tensor_by_name_from_graph(self.graph,
'descrpt_attr/ntypes')
self.rcut = self.descrpt.get_attr('rcut_r')
self.rcut_smth = self.descrpt.get_attr('rcut_r_smth')
self.embedding_net_nodes = get_embedding_net_nodes_from_graph_def(
self.graph_def, suffix=self.suffix)
for tt in self.exclude_types:
if (tt[0] not in range(self.ntypes)) or (tt[1] not in range(
self.ntypes)):
raise RuntimeError("exclude types" + str(tt) +
" must within the number of atomic types " +
str(self.ntypes) + "!")
if (self.ntypes * self.ntypes - len(self.exclude_types) == 0):
raise RuntimeError(
"empty embedding-net are not supported in model compression!")
self.layer_size = len(self.embedding_net_nodes) // (
(self.ntypes * self.ntypes - len(self.exclude_types)) * 2)
self.table_size = self.ntypes * self.ntypes
if type_one_side:
self.layer_size = len(
self.embedding_net_nodes) // (self.ntypes * 2)
self.table_size = self.ntypes
# self.value_type = self.embedding_net_nodes["filter_type_0/matrix_1_0"].dtype #"filter_type_0/matrix_1_0" must exit~
# get trained variables
self.bias = self._get_bias()
self.matrix = self._get_matrix()
for item in self.matrix["layer_" + str(self.layer_size)]:
if len(item) != 0:
self.data_type = type(item[0][0])
self.last_layer_size = item.shape[1]
# define tables
self.data = {}
def init_variables(
self,
graph: tf.Graph,
graph_def: tf.GraphDef,
suffix: str = "",
) -> None:
"""
Init the embedding net variables with the given dict
Parameters
----------
graph : tf.Graph
The input frozen model graph
graph_def : tf.GraphDef
The input frozen model graph_def
suffix : str, optional
The suffix of the scope
"""
super().init_variables(graph=graph, graph_def=graph_def, suffix=suffix)
try:
self.original_sel = get_tensor_by_name_from_graph(
graph, 'descrpt_attr%s/original_sel' % suffix)
except GraphWithoutTensorError:
# original_sel is not restored in old graphs, assume sel never changed before
pass
# check sel == original sel?
try:
sel = get_tensor_by_name_from_graph(graph,
'descrpt_attr%s/sel' % suffix)
except GraphWithoutTensorError:
# sel is not restored in old graphs
pass
else:
if not np.array_equal(np.array(self.sel_a), sel):
if not self.set_davg_zero:
raise RuntimeError(
"Adjusting sel is only supported when `set_davg_zero` is true!"
)
# as set_davg_zero, self.davg is safely zero
self.davg = np.zeros([self.ntypes, self.ndescrpt
]).astype(GLOBAL_NP_FLOAT_PRECISION)
new_dstd = np.ones([self.ntypes, self.ndescrpt
]).astype(GLOBAL_NP_FLOAT_PRECISION)
# shape of davg and dstd is (ntypes, ndescrpt), ndescrpt = 4*sel
n_descpt = np.array(self.sel_a) * 4
n_descpt_old = np.array(sel) * 4
end_index = np.cumsum(n_descpt)
end_index_old = np.cumsum(n_descpt_old)
start_index = np.roll(end_index, 1)
start_index[0] = 0
start_index_old = np.roll(end_index_old, 1)
start_index_old[0] = 0
for nn, oo, ii, jj in zip(n_descpt, n_descpt_old, start_index,
start_index_old):
if nn < oo:
# new size is smaller, copy part of std
new_dstd[:, ii:ii + nn] = self.dstd[:, jj:jj + nn]
else:
# new size is larger, copy all, the rest remains 1
new_dstd[:, ii:ii + oo] = self.dstd[:, jj:jj + oo]
self.dstd = new_dstd
if self.original_sel is None:
self.original_sel = sel
def enable_compression(
self,
min_nbor_dist: float,
model_file: str = 'frozon_model.pb',
table_extrapolate: float = 5,
table_stride_1: float = 0.01,
table_stride_2: float = 0.1,
check_frequency: int = -1,
suffix: str = "",
) -> None:
"""
Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
Parameters
----------
min_nbor_dist
The nearest distance between atoms
model_file
The original frozen model, which will be compressed by the program
table_extrapolate
The scale of model extrapolation
table_stride_1
The uniform stride of the first table
table_stride_2
The uniform stride of the second table
check_frequency
The overflow check frequency
suffix : str, optional
The suffix of the scope
"""
# do some checks before the mocel compression process
assert (
not self.filter_resnet_dt
), "Model compression error: descriptor resnet_dt must be false!"
for tt in self.exclude_types:
if (tt[0] not in range(self.ntypes)) or (tt[1] not in range(
self.ntypes)):
raise RuntimeError("exclude types" + str(tt) +
" must within the number of atomic types " +
str(self.ntypes) + "!")
if (self.ntypes * self.ntypes - len(self.exclude_types) == 0):
raise RuntimeError(
"empty embedding-net are not supported in model compression!")
for ii in range(len(self.filter_neuron) - 1):
if self.filter_neuron[ii] * 2 != self.filter_neuron[ii + 1]:
raise NotImplementedError(
"Model Compression error: descriptor neuron [%s] is not supported by model compression! "
"The size of the next layer of the neural network must be twice the size of the previous layer."
% ','.join([str(item) for item in self.filter_neuron]))
self.compress = True
self.table = DPTabulate(self,
self.filter_neuron,
model_file,
self.type_one_side,
self.exclude_types,
self.compress_activation_fn,
suffix=suffix)
self.table_config = [
table_extrapolate, table_stride_1, table_stride_2, check_frequency
]
self.lower, self.upper \
= self.table.build(min_nbor_dist,
table_extrapolate,
table_stride_1,
table_stride_2)
graph, _ = load_graph_def(model_file)
self.davg = get_tensor_by_name_from_graph(
graph, 'descrpt_attr%s/t_avg' % suffix)
self.dstd = get_tensor_by_name_from_graph(
graph, 'descrpt_attr%s/t_std' % suffix)
def __init__(
self,
descrpt: Descriptor,
neuron: List[int],
model_file: str,
type_one_side: bool = False,
exclude_types: List[List[int]] = [],
activation_fn: Callable[[tf.Tensor], tf.Tensor] = tf.nn.tanh,
suffix: str = "",
) -> None:
"""
Constructor
"""
self.descrpt = descrpt
self.neuron = neuron
self.model_file = model_file
self.type_one_side = type_one_side
self.exclude_types = exclude_types
self.suffix = suffix
# functype
if activation_fn == ACTIVATION_FN_DICT["tanh"]:
self.functype = 1
elif activation_fn == ACTIVATION_FN_DICT["gelu"]:
self.functype = 2
elif activation_fn == ACTIVATION_FN_DICT["relu"]:
self.functype = 3
elif activation_fn == ACTIVATION_FN_DICT["relu6"]:
self.functype = 4
elif activation_fn == ACTIVATION_FN_DICT["softplus"]:
self.functype = 5
elif activation_fn == ACTIVATION_FN_DICT["sigmoid"]:
self.functype = 6
else:
raise RuntimeError("Unknown actication function type!")
self.activation_fn = activation_fn
self.graph, self.graph_def = load_graph_def(self.model_file)
self.sess = tf.Session(graph=self.graph)
self.sub_graph, self.sub_graph_def = self._load_sub_graph()
self.sub_sess = tf.Session(graph=self.sub_graph)
if isinstance(self.descrpt, deepmd.descriptor.DescrptSeR):
try:
self.sel_a = self.graph.get_operation_by_name(
'ProdEnvMatR').get_attr('sel')
self.prod_env_mat_op = self.graph.get_operation_by_name(
'ProdEnvMatR')
except KeyError:
self.sel_a = self.graph.get_operation_by_name(
'DescrptSeR').get_attr('sel')
self.prod_env_mat_op = self.graph.get_operation_by_name(
'DescrptSeR')
elif isinstance(self.descrpt, deepmd.descriptor.DescrptSeA):
try:
self.sel_a = self.graph.get_operation_by_name(
'ProdEnvMatA').get_attr('sel_a')
self.prod_env_mat_op = self.graph.get_operation_by_name(
'ProdEnvMatA')
except KeyError:
self.sel_a = self.graph.get_operation_by_name(
'DescrptSeA').get_attr('sel_a')
self.prod_env_mat_op = self.graph.get_operation_by_name(
'DescrptSeA')
elif isinstance(self.descrpt, deepmd.descriptor.DescrptSeT):
try:
self.sel_a = self.graph.get_operation_by_name(
'ProdEnvMatA').get_attr('sel_a')
self.prod_env_mat_op = self.graph.get_operation_by_name(
'ProdEnvMatA')
except KeyError:
self.sel_a = self.graph.get_operation_by_name(
'DescrptSeA').get_attr('sel_a')
self.prod_env_mat_op = self.graph.get_operation_by_name(
'DescrptSeA')
else:
raise RuntimeError("Unsupported descriptor")
self.davg = get_tensor_by_name_from_graph(
self.graph, f'descrpt_attr{self.suffix}/t_avg')
self.dstd = get_tensor_by_name_from_graph(
self.graph, f'descrpt_attr{self.suffix}/t_std')
self.ntypes = get_tensor_by_name_from_graph(self.graph,
'descrpt_attr/ntypes')
if isinstance(self.descrpt, deepmd.descriptor.DescrptSeR):
self.rcut = self.prod_env_mat_op.get_attr('rcut')
self.rcut_smth = self.prod_env_mat_op.get_attr('rcut_smth')
else:
self.rcut = self.prod_env_mat_op.get_attr('rcut_r')
self.rcut_smth = self.prod_env_mat_op.get_attr('rcut_r_smth')
self.embedding_net_nodes = get_embedding_net_nodes_from_graph_def(
self.graph_def, suffix=self.suffix)
# move it to the descriptor class
# for tt in self.exclude_types:
# if (tt[0] not in range(self.ntypes)) or (tt[1] not in range(self.ntypes)):
# raise RuntimeError("exclude types" + str(tt) + " must within the number of atomic types " + str(self.ntypes) + "!")
# if (self.ntypes * self.ntypes - len(self.exclude_types) == 0):
# raise RuntimeError("empty embedding-net are not supported in model compression!")
self.layer_size = self._get_layer_size()
self.table_size = self._get_table_size()
self.bias = self._get_bias()
self.matrix = self._get_matrix()
self.data_type = self._get_data_type()
self.last_layer_size = self._get_last_layer_size()
self.data = {}