Fix single-task training&data stat

deepmd/pt/model/descriptor/dpa2.py

@@ -304,7 +304,7 @@ def compute_input_stats(self, merged: List[dict], path: Optional[DPPath] = None)
                 }
                 for item in merged
             ]
-            descrpt.compute_input_stats(merged_tmp)
+            descrpt.compute_input_stats(merged_tmp, path)
 
     def serialize(self) -> dict:
         """Serialize the obj to dict."""

deepmd/pt/model/model/__init__.py

@@ -20,7 +20,7 @@
     BaseDescriptor,
 )
 from deepmd.pt.model.task import (
-    Fitting,
+    BaseFitting,
 )
 
 from .dp_model import (
@@ -61,7 +61,7 @@ def get_zbl_model(model_params):
         fitting_net["out_dim"] = descriptor.get_dim_emb()
         if "ener" in fitting_net["type"]:
             fitting_net["return_energy"] = True
-    fitting = Fitting(**fitting_net)
+    fitting = BaseFitting(**fitting_net)
     dp_model = DPAtomicModel(descriptor, fitting, type_map=model_params["type_map"])
     # pairtab
     filepath = model_params["use_srtab"]
@@ -97,9 +97,8 @@ def get_model(model_params):
         fitting_net["out_dim"] = descriptor.get_dim_emb()
         if "ener" in fitting_net["type"]:
             fitting_net["return_energy"] = True
-    fitting = Fitting(**fitting_net)
-
-    model = EnergyModel(descriptor, fitting, type_map=model_params["type_map"])
+    fitting = BaseFitting(**fitting_net)
+    model = DPModel(descriptor, fitting, type_map=model_params["type_map"])
     model.model_def_script = json.dumps(model_params)
     return model
 

deepmd/pt/model/model/dp_model.py

@@ -10,6 +10,7 @@
 )
 from deepmd.pt.model.task.ener import (
     EnergyFittingNet,
+    EnergyFittingNetDirect,
 )
 from deepmd.pt.model.task.polarizability import (
     PolarFittingNet,
@@ -36,7 +37,9 @@ def __new__(cls, descriptor, fitting, *args, **kwargs):
         # according to the fitting network to decide the type of the model
         if cls is DPModel:
             # map fitting to model
-            if isinstance(fitting, EnergyFittingNet):
+            if isinstance(fitting, EnergyFittingNet) or isinstance(
+                fitting, EnergyFittingNetDirect
+            ):
                 cls = EnergyModel
             elif isinstance(fitting, DipoleFittingNet):
                 cls = DipoleModel

deepmd/pt/model/model/model.py

@@ -59,9 +59,9 @@
 # in DPAtomicModel (and other classes), but this requires the developer aware
 # of it when developing it...
 class BaseModel(make_base_model()):
-    def __init__(self):
+    def __init__(self, *args, **kwargs):
         """Construct a basic model for different tasks."""
-        super().__init__()
+        super().__init__(*args, **kwargs)
 
     def compute_or_load_stat(
         self,

deepmd/utils/path.py

@@ -355,6 +355,7 @@
         if self._name in self._keys:
             del self.root[self._name]
         self.root.create_dataset(self._name, data=arr)
+        self.root.flush()
 
     def glob(self, pattern: str) -> List["DPPath"]:
         """Search path using the glob pattern.

examples/water/dpa2/input_torch.json

@@ -1,18 +1,13 @@
 {
   "_comment": "that's all",
   "model": {
-    "type_embedding": {
-      "neuron": [
-        8
-      ],
-      "tebd_input_mode": "concat"
-    },
     "type_map": [
       "O",
       "H"
     ],
     "descriptor": {
       "type": "dpa2",
+      "tebd_dim": 8,
       "repinit_rcut": 9.0,
       "repinit_rcut_smth": 8.0,
       "repinit_nsel": 120,
@@ -74,6 +69,7 @@
     "_comment": " that's all"
   },
   "training": {
+    "stat_file": "./dpa2",
     "training_data": {
       "systems": [
         "../data/data_0",

examples/water/se_atten/input_torch.json

@@ -15,6 +15,7 @@
         50,
         100
       ],
+      "tebd_dim": 8,
       "axis_neuron": 16,
       "attn": 128,
       "attn_layer": 2,
@@ -59,6 +60,7 @@
     "_comment": " that's all"
   },
   "training": {
+    "stat_file": "./dpa1",
     "training_data": {
       "systems": [
         "../data/data_0",

examples/water/se_e2_a/input_torch.json

@@ -51,6 +51,7 @@
     "_comment": " that's all"
   },
   "training": {
+    "stat_file": "./se_e2_a",
     "training_data": {
       "systems": [
         "../data/data_0",