Merge c8ded77 into 843a3c5

deepmodeling · Nov 26, 2021 · 2111e46 · 2111e46
2 parents 843a3c5 + c8ded77
commit 2111e46
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 2 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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]
                         }