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 = {}