diff --git a/deepmd/entrypoints/main.py b/deepmd/entrypoints/main.py index 5546ca15cd..657dff2bbb 100644 --- a/deepmd/entrypoints/main.py +++ b/deepmd/entrypoints/main.py @@ -445,6 +445,12 @@ def main_parser() -> argparse.ArgumentParser: required=True, help="type map", ) + parser_neighbor_stat.add_argument( + "--one-type", + action="store_true", + default=False, + help="treat all types as a single type. Used with se_atten descriptor.", + ) # --version parser.add_argument('--version', action='version', version='DeePMD-kit v%s' % __version__) diff --git a/deepmd/entrypoints/neighbor_stat.py b/deepmd/entrypoints/neighbor_stat.py index eaf8b3efe3..bc3c6430ff 100644 --- a/deepmd/entrypoints/neighbor_stat.py +++ b/deepmd/entrypoints/neighbor_stat.py @@ -12,6 +12,7 @@ def neighbor_stat( system: str, rcut: float, type_map: List[str], + one_type: bool = False, **kwargs, ): """Calculate neighbor statistics. @@ -24,6 +25,8 @@ def neighbor_stat( cutoff radius type_map : list[str] type map + one_type : bool, optional, default=False + treat all types as a single type Examples -------- @@ -42,7 +45,7 @@ def neighbor_stat( type_map=type_map, ) data.get_batch() - nei = NeighborStat(data.get_ntypes(), rcut) + nei = NeighborStat(data.get_ntypes(), rcut, one_type=one_type) min_nbor_dist, max_nbor_size = nei.get_stat(data) log.info("min_nbor_dist: %f" % min_nbor_dist) log.info("max_nbor_size: %s" % str(max_nbor_size)) diff --git a/deepmd/entrypoints/train.py b/deepmd/entrypoints/train.py index 7cdb0307d4..7ffdc7c8b0 100755 --- a/deepmd/entrypoints/train.py +++ b/deepmd/entrypoints/train.py @@ -90,10 +90,7 @@ def train( jdata = normalize(jdata) - if jdata['model']['descriptor']['type'] in ['se_atten'] and isinstance(jdata['model']['descriptor']['sel'], list): - jdata['model']['descriptor']['sel'] = sum(jdata['model']['descriptor']['sel']) - - if not is_compress and not skip_neighbor_stat and jdata['model']['descriptor']['type'] not in ['se_atten']: + if not is_compress and not skip_neighbor_stat: jdata = update_sel(jdata) with open(output, "w") as fp: @@ -253,7 +250,7 @@ def get_type_map(jdata): return jdata['model'].get('type_map', None) -def get_nbor_stat(jdata, rcut): +def get_nbor_stat(jdata, rcut, one_type: bool = False): max_rcut = get_rcut(jdata) type_map = get_type_map(jdata) @@ -268,7 +265,7 @@ def get_nbor_stat(jdata, rcut): map_ntypes = data_ntypes ntypes = max([map_ntypes, data_ntypes]) - neistat = NeighborStat(ntypes, rcut) + neistat = NeighborStat(ntypes, rcut, one_type=one_type) min_nbor_dist, max_nbor_size = neistat.get_stat(train_data) @@ -283,8 +280,8 @@ def get_nbor_stat(jdata, rcut): dtype = tf.int32) return min_nbor_dist, max_nbor_size -def get_sel(jdata, rcut): - _, max_nbor_size = get_nbor_stat(jdata, rcut) +def get_sel(jdata, rcut, one_type: bool = False): + _, max_nbor_size = get_nbor_stat(jdata, rcut, one_type=one_type) return max_nbor_size def get_min_nbor_dist(jdata, rcut): @@ -320,14 +317,20 @@ def wrap_up_4(xx): def update_one_sel(jdata, descriptor): + if descriptor['type'] == 'loc_frame': + return descriptor rcut = descriptor['rcut'] - tmp_sel = get_sel(jdata, rcut) + tmp_sel = get_sel(jdata, rcut, one_type=descriptor['type'] in ('se_atten',)) + sel = descriptor['sel'] + if isinstance(sel, int): + # convert to list and finnally convert back to int + sel = [sel] if parse_auto_sel(descriptor['sel']) : ratio = parse_auto_sel_ratio(descriptor['sel']) - descriptor['sel'] = [int(wrap_up_4(ii * ratio)) for ii in tmp_sel] + descriptor['sel'] = sel = [int(wrap_up_4(ii * ratio)) for ii in tmp_sel] else: # sel is set by user - for ii, (tt, dd) in enumerate(zip(tmp_sel, descriptor['sel'])): + for ii, (tt, dd) in enumerate(zip(tmp_sel, sel)): if dd and tt > dd: # we may skip warning for sel=0, where the user is likely # to exclude such type in the descriptor @@ -336,6 +339,8 @@ def update_one_sel(jdata, descriptor): "not less than %d, but you set it to %d. The accuracy" " of your model may get worse." %(ii, tt, dd) ) + if descriptor['type'] in ('se_atten',): + descriptor['sel'] = sel = sum(sel) return descriptor @@ -344,9 +349,8 @@ def update_sel(jdata): descrpt_data = jdata['model']['descriptor'] if descrpt_data['type'] == 'hybrid': for ii in range(len(descrpt_data['list'])): - if descrpt_data['list'][ii]['type'] != 'loc_frame': - descrpt_data['list'][ii] = update_one_sel(jdata, descrpt_data['list'][ii]) - elif descrpt_data['type'] != 'loc_frame': + descrpt_data['list'][ii] = update_one_sel(jdata, descrpt_data['list'][ii]) + else: descrpt_data = update_one_sel(jdata, descrpt_data) jdata['model']['descriptor'] = descrpt_data return jdata diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 7ba62f9953..1f84b7fe60 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -243,7 +243,8 @@ def descrpt_hybrid_args(): def descrpt_se_atten_args(): doc_sel = 'This parameter set the number of selected neighbors. Note that this parameter is a little different from that in other descriptors. Instead of separating each type of atoms, only the summation matters. And this number is highly related with the efficiency, thus one should not make it too large. Usually 200 or less is enough, far away from the GPU limitation 4096. It can be:\n\n\ - `int`. The maximum number of neighbor atoms to be considered. We recommend it to be less than 200. \n\n\ - - `List[int]`. The length of the list should be the same as the number of atom types in the system. `sel[i]` gives the selected number of type-i neighbors. Only the summation of `sel[i]` matters, and it is recommended to be less than 200.' + - `List[int]`. The length of the list should be the same as the number of atom types in the system. `sel[i]` gives the selected number of type-i neighbors. Only the summation of `sel[i]` matters, and it is recommended to be less than 200.\ + - `str`. Can be "auto:factor" or "auto". "factor" is a float number larger than 1. This option will automatically determine the `sel`. In detail it counts the maximal number of neighbors with in the cutoff radius for each type of neighbor, then multiply the maximum by the "factor". Finally the number is wraped up to 4 divisible. The option "auto" is equivalent to "auto:1.1".' doc_rcut = 'The cut-off radius.' doc_rcut_smth = 'Where to start smoothing. For example the 1/r term is smoothed from `rcut` to `rcut_smth`' doc_neuron = 'Number of neurons in each hidden layers of the embedding net. When two layers are of the same size or one layer is twice as large as the previous layer, a skip connection is built.' @@ -262,7 +263,7 @@ def descrpt_se_atten_args(): doc_attn_mask = 'Whether to do mask on the diagonal in the attention matrix' return [ - Argument("sel", [int, list], optional=True, default=200, doc=doc_sel), + Argument("sel", [int, list, str], optional=True, default="auto", doc=doc_sel), Argument("rcut", float, optional=True, default=6.0, doc=doc_rcut), Argument("rcut_smth", float, optional=True, default=0.5, doc=doc_rcut_smth), Argument("neuron", list, optional=True, default=[10, 20, 40], doc=doc_neuron), diff --git a/deepmd/utils/neighbor_stat.py b/deepmd/utils/neighbor_stat.py index 9088cfc8a4..87e5b83774 100644 --- a/deepmd/utils/neighbor_stat.py +++ b/deepmd/utils/neighbor_stat.py @@ -23,15 +23,20 @@ class NeighborStat(): The num of atom types rcut The cut-off radius + one_type : bool, optional, default=False + Treat all types as a single type. """ def __init__(self, ntypes : int, - rcut: float) -> None: + rcut: float, + one_type : bool = False, + ) -> None: """ Constructor """ self.rcut = rcut self.ntypes = ntypes + self.one_type = one_type sub_graph = tf.Graph() def builder(): @@ -41,10 +46,17 @@ def builder(): place_holders['type'] = tf.placeholder(tf.int32, [None, None], name='t_type') place_holders['natoms_vec'] = tf.placeholder(tf.int32, [self.ntypes+2], name='t_natoms') place_holders['default_mesh'] = tf.placeholder(tf.int32, [None], name='t_mesh') + t_type = place_holders['type'] + t_natoms = place_holders['natoms_vec'] + if self.one_type: + # all types = 0, natoms_vec = [natoms, natoms, natoms] + t_type = tf.zeros_like(t_type, dtype=tf.int32) + t_natoms = tf.repeat(t_natoms[0], 3) + _max_nbor_size, _min_nbor_dist \ = op_module.neighbor_stat(place_holders['coord'], - place_holders['type'], - place_holders['natoms_vec'], + t_type, + t_natoms, place_holders['box'], place_holders['default_mesh'], rcut = self.rcut) @@ -74,7 +86,9 @@ def get_stat(self, A list with ntypes integers, denotes the actual achieved max sel """ self.min_nbor_dist = 100.0 - self.max_nbor_size = [0] * self.ntypes + self.max_nbor_size = [0] + if not self.one_type: + self.max_nbor_size *= self.ntypes def feed(): for ii in range(len(data.system_dirs)): diff --git a/source/tests/test_train.py b/source/tests/test_train.py index 2b7fd16d18..ade7ab7cbd 100644 --- a/source/tests/test_train.py +++ b/source/tests/test_train.py @@ -29,6 +29,7 @@ def test_update_one_sel(self, sel_mock): sel_mock.return_value = [10,20] jdata = {} descriptor = { + 'type': 'se_e2_a', 'rcut': 6, 'sel': "auto" } @@ -36,6 +37,7 @@ def test_update_one_sel(self, sel_mock): # self.assertEqual(descriptor['sel'], [11,22]) self.assertEqual(descriptor['sel'], [12,24]) descriptor = { + 'type': 'se_e2_a', 'rcut': 6, 'sel': "auto:1.5" } @@ -113,7 +115,98 @@ def test_update_sel(self, sel_mock): jdata = update_sel(jdata) self.assertEqual(jdata, expected_out) - + @patch("deepmd.entrypoints.train.get_sel") + def test_update_sel_atten_auto(self, sel_mock): + sel_mock.return_value = [25] + jdata = { + 'model' : { + 'descriptor': { + 'type' : 'se_atten', + 'sel' : "auto", + 'rcut': 6, + } + } + } + expected_out = { + 'model' : { + 'descriptor': { + 'type' : 'se_atten', + 'sel' : 28, + 'rcut': 6, + } + } + } + jdata = update_sel(jdata) + self.assertEqual(jdata, expected_out) + + @patch("deepmd.entrypoints.train.get_sel") + def test_update_sel_atten_int(self, sel_mock): + sel_mock.return_value = [25] + jdata = { + 'model' : { + 'descriptor': { + 'type' : 'se_atten', + 'sel' : 30, + 'rcut': 6, + } + } + } + expected_out = { + 'model' : { + 'descriptor': { + 'type' : 'se_atten', + 'sel' : 30, + 'rcut': 6, + } + } + } + jdata = update_sel(jdata) + self.assertEqual(jdata, expected_out) + + @patch("deepmd.entrypoints.train.get_sel") + def test_update_sel_atten_list(self, sel_mock): + sel_mock.return_value = [25] + jdata = { + 'model' : { + 'descriptor': { + 'type' : 'se_atten', + 'sel' : 30, + 'rcut': 6, + } + } + } + expected_out = { + 'model' : { + 'descriptor': { + 'type' : 'se_atten', + 'sel' : 30, + 'rcut': 6, + } + } + } + jdata = update_sel(jdata) + self.assertEqual(jdata, expected_out) + + def test_skip_loc_frame(self): + jdata = { + 'model' : { + 'descriptor': { + 'type' : 'loc_frame', + 'rcut': 6, + } + } + } + expected_out = { + 'model' : { + 'descriptor': { + 'type' : 'loc_frame', + 'rcut': 6, + } + } + } + jdata = update_sel(jdata) + self.assertEqual(jdata, expected_out) + def test_wrap_up_4(self): self.assertEqual(wrap_up_4(12), 3 * 4) self.assertEqual(wrap_up_4(13), 4 * 4)