def _init_from_frz_model(self):
# get the model type from the frozen model(self.run_opt.init_frz_model)
try:
t_model_type = get_tensor_by_name(self.run_opt.init_frz_model, 'model_type')
self.model_type = bytes.decode(t_model_type)
except GraphWithoutTensorError as e:
# throw runtime error if there's no frozen model
if not os.path.exists(self.run_opt.init_frz_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
# throw runtime error if the frozen_model has no model type information...
else:
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
if self.fitting_type != 'ener':
raise RuntimeError("The 'dp train init-frz-model' command only supports the 'ener' type fitting net currently!")
# self.frz_model will control the self.model to import the descriptor from the given frozen model instead of building from scratch...
# 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))
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))
tf.constant("compressed_model", name = 'model_type', dtype = tf.string)
else:
raise RuntimeError("Unknown model type %s" % self.model_type)
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)
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()
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:
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']))
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)
self._build_lr()
self._build_network(data)
self._build_training()