update NvNMD training code

deepmd/main.py

@@ -778,6 +778,13 @@ def main_parser() -> argparse.ArgumentParser:
         default=None,
         help="Restart the training from the provided prefix of checkpoint files.",
     )
+    parser_train_nvnmd.add_argument(
+        "-f",
+        "--init-frz-model",
+        type=str,
+        default=None,
+        help="Initialize the training from the frozen model.",
+    )
     parser_train_nvnmd.add_argument(
         "-s",
         "--step",

deepmd/tf/descriptor/se_atten.py

@@ -38,8 +38,8 @@
 from deepmd.tf.nvnmd.descriptor.se_atten import (
     build_davg_dstd,
     build_op_descriptor,
+    build_recovered,
     check_switch_range,
-    descrpt2r4,
     filter_GR2D,
     filter_lower_R42GR,
 )
@@ -715,6 +715,20 @@ def _pass_filter(
         inputs_i = inputs
         inputs_i = tf.reshape(inputs_i, [-1, self.ndescrpt])
         type_i = -1
+
+        # descrpt and recovered_switch for nvnmd
+        if nvnmd_cfg.enable and nvnmd_cfg.quantize_descriptor:
+            inputs_i, self.recovered_switch = build_recovered(
+                inputs_i,
+                self.t_avg,
+                self.t_std,
+                self.atype_nloc,
+                natoms[0],
+                self.ntypes,
+                self.rcut_r_smth,
+                self.filter_precision,
+            )
+
         if len(self.exclude_types):
             mask = self.build_type_exclude_mask_mixed(
                 self.exclude_types,
@@ -753,8 +767,7 @@ def _pass_filter(
                 )
                 self.negative_mask = -(2 << 32) * (1.0 - self.nmask)
                 inputs_i *= mask
-        if nvnmd_cfg.enable and nvnmd_cfg.quantize_descriptor:
-            inputs_i = descrpt2r4(inputs_i, atype)
+
         layer, qmat = self._filter(
             inputs_i,
             type_i,
@@ -1162,6 +1175,7 @@ def _filter_lower(
                                 inputs_i,
                                 atype,
                                 self.nei_type_vec,
+                                self.recovered_switch,
                             )
                         elif nvnmd_cfg.restore_descriptor:
                             self.embedding_net_variables = (

deepmd/tf/nvnmd/data/data.py

@@ -1,3 +1,5 @@
+import copy
+
 # SPDX-License-Identifier: LGPL-3.0-or-later
 jdata_sys = {"debug": False}
 
@@ -10,6 +12,7 @@
         "neuron": [8, 16, 32],
         "resnet_dt": False,
         "axis_neuron": 4,
+        "seed": 1,
         "type_one_side": True,
         # rcut range
         "rc_lim": 0.5,
@@ -39,6 +42,7 @@
         # basic config from deepmd model
         "neuron": [128, 128, 128],
         "resnet_dt": False,
+        "seed": 1,
         "NNODE_FITS": "(M1*M2, neuron, 1)",
         "nlayer_fit": "len(neuron)+1",
         "NLAYER": "nlayer_fit",
@@ -118,9 +122,9 @@
     "end": "",
 }
 
-# change the configuration according to the max_nnei
-jdata_config_v0_ni128 = jdata_config_v0.copy()
-jdata_config_v0_ni256 = jdata_config_v0.copy()
+# change the configuration accordng to the max_nnei
+jdata_config_v0_ni128 = copy.deepcopy(jdata_config_v0)
+jdata_config_v0_ni256 = copy.deepcopy(jdata_config_v0)
 jdata_config_v0_ni256["ctrl"] = {
     "MAX_NNEI": 256,
     "NSTDM": 128,
@@ -142,6 +146,7 @@
         "rcut_smth": 0.5,
         "neuron": [8, 16, 32],
         "resnet_dt": False,
+        "seed": 1,
         "axis_neuron": 4,
         "type_one_side": True,
         # rcut range
@@ -150,6 +155,7 @@
         # embedding net size
         "M1": "neuron[-1]",
         "M2": "axis_neuron",
+        "M3": 2,
         "SEL": 128,
         "NNODE_FEAS": "(1, neuron)",
         "nlayer_fea": "len(neuron)",
@@ -170,6 +176,7 @@
         # basic config from deepmd model
         "neuron": [128, 128, 128],
         "resnet_dt": False,
+        "seed": 1,
         "NNODE_FITS": "(M1*M2, neuron, 1)",
         "nlayer_fit": "len(neuron)+1",
         "NLAYER": "nlayer_fit",
@@ -200,7 +207,7 @@
         "NSEL": "NSTDM",
         "NSADV": "NSTDM+1",
         "VERSION": 1,
-        "SUB_VERSION": 1,
+        "SUB_VERSION": 2,
     },
     "nbit": {
         # general
@@ -243,26 +250,27 @@
         "NBIT_CFG": 64,
         "NBIT_NET": 72,
         "NBIT_MODEL_HEAD": 32,
+        "NBIT_NSTEP": 8,
         # nbit for mapt-version
         "NBIT_IDX_S2G": 9,
         "NBIT_NEIB": 8,
     },
     "end": "",
 }
 
-# change the configuration according to the max_nnei
-jdata_config_v1_ni128 = jdata_config_v1.copy()
-jdata_config_v1_ni256 = jdata_config_v1.copy()
+# change the configuration accordng to the max_nnei
+jdata_config_v1_ni128 = copy.deepcopy(jdata_config_v1)
+jdata_config_v1_ni256 = copy.deepcopy(jdata_config_v1)
 jdata_config_v1_ni256["ctrl"] = {
     "MAX_NNEI": 256,
-    "NSTDM": 128,
+    "NSTDM": 64,
     "NSTDM_M1": 32,
-    "NSTDM_M2": 4,
+    "NSTDM_M2": 2,
     "NSTDM_M1X": 8,
     "NSEL": "NSTDM",
     "NSADV": "NSTDM+1",
     "VERSION": 1,
-    "SUB_VERSION": 1,
+    "SUB_VERSION": 2,
 }
 jdata_config_v1_ni256["nbit"]["NBIT_NEIB"] = 9
 
@@ -283,6 +291,7 @@
     },
     "nvnmd": {
         "version": 0,
+        "device": "vu9p",
         "max_nnei": 128,  # 128 or 256
         "net_size": 128,
         "config_file": "none",
@@ -319,20 +328,22 @@
         "disp_training": True,
         "time_training": True,
         "profiling": False,
-        "training_data": {"systems": "dataset", "batch_size": 1},
+        "training_data": {"systems": "dataset", "set_prefix": "set", "batch_size": 1},
     },
 }
 
-jdata_deepmd_input_v0_ni128 = jdata_deepmd_input_v0.copy()
-jdata_deepmd_input_v0_ni256 = jdata_deepmd_input_v0.copy()
+jdata_deepmd_input_v0_ni128 = copy.deepcopy(jdata_deepmd_input_v0)
+jdata_deepmd_input_v0_ni256 = copy.deepcopy(jdata_deepmd_input_v0)
 jdata_deepmd_input_v0_ni256["nvnmd"]["max_nnei"] = 256
 
 jdata_deepmd_input_v1 = {
     "model": {
         "descriptor": {
             "seed": 1,
             "type": "se_atten",
-            "tebd_input_mode": "strip",
+            "stripped_type_embedding": True,
+            "smooth_type_embdding": True,
+            "set_davg_zero": False,
             "sel": 128,
             "rcut": 7.0,
             "rcut_smth": 0.5,
@@ -349,6 +360,7 @@
     },
     "nvnmd": {
         "version": 1,
+        "device": "vu9p",
         "max_nnei": 128,  # 128 or 256
         "net_size": 128,
         "config_file": "none",
@@ -385,14 +397,15 @@
         "disp_training": True,
         "time_training": True,
         "profiling": False,
-        "training_data": {"systems": "dataset", "batch_size": 1},
+        "training_data": {"systems": "dataset", "set_prefix": "set", "batch_size": 1},
     },
 }
 
-jdata_deepmd_input_v1_ni128 = jdata_deepmd_input_v1.copy()
-jdata_deepmd_input_v1_ni256 = jdata_deepmd_input_v1.copy()
+jdata_deepmd_input_v1_ni128 = copy.deepcopy(jdata_deepmd_input_v1)
+jdata_deepmd_input_v1_ni256 = copy.deepcopy(jdata_deepmd_input_v1)
 jdata_deepmd_input_v1_ni256["nvnmd"]["max_nnei"] = 256
 
+
 NVNMD_WELCOME = (
     r" _   _  __     __  _   _   __  __   ____  ",
     r"| \ | | \ \   / / | \ | | |  \/  | |  _ \ ",