diff --git a/deepmd/entrypoints/main.py b/deepmd/entrypoints/main.py index 4f772526c8..632d35fd71 100644 --- a/deepmd/entrypoints/main.py +++ b/deepmd/entrypoints/main.py @@ -170,6 +170,11 @@ def parse_args(args: Optional[List[str]] = None): default=None, help="Initialize the training from the frozen model.", ) + parser_train.add_argument( + "--skip-neighbor-stat", + action="store_true", + help="Skip calculating neighbor statistics. Sel checking, automatic sel, and model compression will be disabled.", + ) # * freeze script ****************************************************************** parser_frz = subparsers.add_parser( diff --git a/deepmd/entrypoints/train.py b/deepmd/entrypoints/train.py index 093c95fbff..42e21d2ba4 100755 --- a/deepmd/entrypoints/train.py +++ b/deepmd/entrypoints/train.py @@ -38,6 +38,7 @@ def train( log_level: int, log_path: Optional[str], is_compress: bool = False, + skip_neighbor_stat: bool = False, **kwargs, ): """Run DeePMD model training. @@ -62,6 +63,8 @@ def train( logging file path or None if logs are to be output only to stdout is_compress: bool indicates whether in the model compress mode + skip_neighbor_stat : bool, default=False + skip checking neighbor statistics Raises ------ @@ -87,7 +90,7 @@ def train( jdata = normalize(jdata) - if not is_compress: + if not is_compress and not skip_neighbor_stat: jdata = update_sel(jdata) with open(output, "w") as fp: @@ -333,10 +336,12 @@ def update_one_sel(jdata, descriptor): def update_sel(jdata): + log.info("Calculate neighbor statistics... (add --skip-neighbor-stat to skip this step)") descrpt_data = jdata['model']['descriptor'] if descrpt_data['type'] == 'hybrid': for ii in range(len(descrpt_data['list'])): - descrpt_data['list'][ii] = update_one_sel(jdata, descrpt_data['list'][ii]) + 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 = update_one_sel(jdata, descrpt_data) jdata['model']['descriptor'] = descrpt_data diff --git a/doc/train/training-advanced.md b/doc/train/training-advanced.md index 865b12e460..7b4c475ae7 100644 --- a/doc/train/training-advanced.md +++ b/doc/train/training-advanced.md @@ -114,6 +114,7 @@ optional arguments: --init-frz-model INIT_FRZ_MODEL Initialize the training from the frozen model. + --skip-neighbor-stat Skip calculating neighbor statistics. Sel checking, automatic sel, and model compression will be disabled. (default: False) ``` **`--init-model model.ckpt`**, initializes the model training with an existing model that is stored in the checkpoint `model.ckpt`, the network architectures should match. @@ -122,6 +123,8 @@ optional arguments: **`--init-frz-model frozen_model.pb`**, initializes the training with an existing model that is stored in `frozen_model.pb`. +**`--skip-neighbor-stat`** will skip calculating neighbor statistics if one is concerned about performance. Some features will be disabled. + To get the best performance, one should control the number of threads used by DeePMD-kit. This is achieved by three environmental variables: `OMP_NUM_THREADS`, `TF_INTRA_OP_PARALLELISM_THREADS` and `TF_INTER_OP_PARALLELISM_THREADS`. `OMP_NUM_THREADS` controls the multithreading of DeePMD-kit implemented operations. `TF_INTRA_OP_PARALLELISM_THREADS` and `TF_INTER_OP_PARALLELISM_THREADS` controls `intra_op_parallelism_threads` and `inter_op_parallelism_threads`, which are Tensorflow configurations for multithreading. An explanation is found [here](https://www.intel.com/content/www/us/en/developer/articles/technical/maximize-tensorflow-performance-on-cpu-considerations-and-recommendations-for-inference.html). For example if you wish to use 3 cores of 2 CPUs on one node, you may set the environmental variables and run DeePMD-kit as follows: diff --git a/source/tests/test_train.py b/source/tests/test_train.py index 21892fb8fe..2b7fd16d18 100644 --- a/source/tests/test_train.py +++ b/source/tests/test_train.py @@ -53,10 +53,12 @@ def test_update_sel_hybrid(self, sel_mock): 'type' : 'hybrid', 'list' : [ { + 'type': 'se_e2_a', 'rcut': 6, 'sel': "auto" }, { + 'type': 'se_e2_a', 'rcut': 6, 'sel': "auto:1.5" } @@ -70,10 +72,12 @@ def test_update_sel_hybrid(self, sel_mock): 'type' : 'hybrid', 'list' : [ { + 'type': 'se_e2_a', 'rcut': 6, 'sel': [12,24] }, { + 'type': 'se_e2_a', 'rcut': 6, 'sel': [16,32] }