示例#1
0
 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="")
示例#2
0
    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)
示例#3
0
 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)
示例#4
0
文件: ener.py 项目: njzjz/deepmd-kit
    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)
示例#5
0
    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 = {}
示例#6
0
文件: se_a.py 项目: njzjz/deepmd-kit
    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)
示例#7
0
    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 = {}
示例#8
0
    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()