def _import_graph_def_from_frz_model(self, frz_model, feed_dict,
return_elements):
graph, graph_def = load_graph_def(frz_model)
return tf.import_graph_def(graph_def,
input_map=feed_dict,
return_elements=return_elements,
name="")
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_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 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__(
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 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 = {}
def build(self, data=None, stop_batch=0):
self.ntypes = self.model.get_ntypes()
self.stop_batch = stop_batch
if self.numb_fparam > 0:
log.info("training with %d frame parameter(s)" % self.numb_fparam)
else:
log.info("training without frame parameter")
if not self.is_compress:
# Usually, the type number of the model should be equal to that of the data
# However, nt_model > nt_data should be allowed, since users may only want to
# train using a dataset that only have some of elements
if self.ntypes < data.get_ntypes():
raise ValueError(
"The number of types of the training data is %d, but that of the "
"model is only %d. The latter must be no less than the former. "
"You may need to reset one or both of them. Usually, the former "
"is given by `model/type_map` in the training parameter (if set) "
"or the maximum number in the training data. The latter is given "
"by `model/descriptor/sel` in the training parameter." %
(data.get_ntypes(), self.ntypes))
self.type_map = data.get_type_map()
self.batch_size = data.get_batch_size()
if self.run_opt.init_mode not in ('init_from_model', 'restart',
'init_from_frz_model'):
# self.saver.restore (in self._init_session) will restore avg and std variables, so data_stat is useless
# init_from_frz_model will restore data_stat variables in `init_variables` method
log.info("data stating... (this step may take long time)")
self.model.data_stat(data)
# config the init_frz_model command
if self.run_opt.init_mode == 'init_from_frz_model':
self._init_from_frz_model()
# neighbor_stat is moved to train.py as duplicated
# TODO: this is a simple fix but we should have a clear
# architecture to call neighbor stat
else:
graph, graph_def = load_graph_def(
self.model_param['compress']['model_file'])
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(graph, graph_def)
# 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)
else:
tf.constant("original_model", name='model_type', dtype=tf.string)
if self.mixed_prec is not None:
self.descrpt.enable_mixed_precision(self.mixed_prec)
self.fitting.enable_mixed_precision(self.mixed_prec)
self._build_lr()
self._build_network(data)
self._build_training()