pt: refactor data stat

deepmd/pt/model/atomic_model/dp_atomic_model.py

@@ -22,7 +22,7 @@
 from deepmd.pt.utils.utils import (
     dict_to_device,
 )
 from deepmd.utils.path import (
     DPPath,
 )
 
@@ -179,13 +179,17 @@
         stat_file_path
             The dictionary of paths to the statistics files.
         """
+        if stat_file_path is not None and self.type_map is not None:
+            # descriptors and fitting net with different type_map
+            # should not share the same parameters
+            stat_file_path /= " ".join(self.type_map)
         for data_sys in sampled:
             dict_to_device(data_sys)
         if sampled is None:
             sampled = []
         self.descriptor.compute_input_stats(sampled, stat_file_path)
         if self.fitting_net is not None:
             self.fitting_net.compute_output_stats(sampled, stat_file_path)
 
     @torch.jit.export
     def get_dim_fparam(self) -> int:

deepmd/pt/model/descriptor/descriptor.py

@@ -15,10 +15,19 @@
 from deepmd.pt.model.network.network import (
     TypeEmbedNet,
 )
+from deepmd.pt.utils import (
+    env,
+)
+from deepmd.pt.utils.env_mat_stat import (
+    EnvMatStatSeA,
+)
 from deepmd.pt.utils.plugin import (
     Plugin,
 )
+from deepmd.utils.env_mat_stat import (
+    StatItem,
+)
 from deepmd.utils.path import (
     DPPath,
 )
 
@@ -171,10 +180,14 @@
         """Returns the embedding dimension."""
         pass
 
     def compute_input_stats(self, merged: List[dict], path: Optional[DPPath] = None):
         """Update mean and stddev for DescriptorBlock elements."""
         raise NotImplementedError
 
+    def get_stats(self) -> dict[str, StatItem]:
+        """Get the statistics of the descriptor."""
+        raise NotImplementedError
+
     def share_params(self, base_class, shared_level, resume=False):
         assert (
             self.__class__ == base_class.__class__
@@ -183,27 +196,13 @@
             # link buffers
             if hasattr(self, "mean") and not resume:
                 # in case of change params during resume
-                sumr_base, suma_base, sumn_base, sumr2_base, suma2_base = (
-                    base_class.sumr,
-                    base_class.suma,
-                    base_class.sumn,
-                    base_class.sumr2,
-                    base_class.suma2,
-                )
-                sumr, suma, sumn, sumr2, suma2 = (
-                    self.sumr,
-                    self.suma,
-                    self.sumn,
-                    self.sumr2,
-                    self.suma2,
-                )
-                stat_dict = {
-                    "sumr": sumr_base + sumr,
-                    "suma": suma_base + suma,
-                    "sumn": sumn_base + sumn,
-                    "sumr2": sumr2_base + sumr2,
-                    "suma2": suma2_base + suma2,
-                }
+                base_env = EnvMatStatSeA(base_class)
+                for kk in base_class.get_stats():
+                    base_env.stats[kk] += self.get_stats()[kk]
+                mean, stddev = base_env()
+                if not base_class.set_davg_zero:
+                    base_class.mean.copy_(torch.tensor(mean, device=env.DEVICE))
+                base_class.stddev.copy_(torch.tensor(stddev, device=env.DEVICE))
                 self.mean = base_class.mean
                 self.stddev = base_class.stddev
             # self.load_state_dict(base_class.state_dict()) # this does not work, because it only inits the model

deepmd/pt/model/descriptor/repformers.py

@@ -18,13 +18,16 @@
 from deepmd.pt.utils import (
     env,
 )
 from deepmd.pt.utils.env_mat_stat import (
     EnvMatStatSeA,
 )
 from deepmd.pt.utils.utils import (
     get_activation_fn,
 )
+from deepmd.utils.env_mat_stat import (
+    StatItem,
+)
 from deepmd.utils.path import (
     DPPath,
 )
 
@@ -147,6 +150,7 @@
         stddev = torch.ones(sshape, dtype=mydtype, device=mydev)
         self.register_buffer("mean", mean)
         self.register_buffer("stddev", stddev)
+        self.stats = None
 
     def get_rcut(self) -> float:
         """Returns the cut-off radius."""
@@ -266,13 +270,22 @@
 
         return g1, g2, h2, rot_mat.view(-1, nloc, self.dim_emb, 3), sw
 
     def compute_input_stats(self, merged: List[dict], path: Optional[DPPath] = None):
         """Update mean and stddev for descriptor elements."""
         env_mat_stat = EnvMatStatSeA(self)
         if path is not None:
             path = path / env_mat_stat.get_hash()
         env_mat_stat.load_or_compute_stats(merged, path)
+        self.stats = env_mat_stat.stats
         mean, stddev = env_mat_stat()
         if not self.set_davg_zero:
             self.mean.copy_(torch.tensor(mean, device=env.DEVICE))
         self.stddev.copy_(torch.tensor(stddev, device=env.DEVICE))
+
+    def get_stats(self) -> dict[str, StatItem]:
+        """Get the statistics of the descriptor."""
+        if self.stats is None:
+            raise RuntimeError(
+                "The statistics of the descriptor has not been computed."
+            )
+        return self.stats

deepmd/pt/model/descriptor/se_a.py

@@ -21,10 +21,13 @@
     PRECISION_DICT,
     RESERVED_PRECISON_DICT,
 )
 from deepmd.pt.utils.env_mat_stat import (
     EnvMatStatSeA,
 )
+from deepmd.utils.env_mat_stat import (
+    StatItem,
+)
 from deepmd.utils.path import (
     DPPath,
 )
 
@@ -116,9 +119,9 @@
         """Returns the output dimension of this descriptor."""
         return self.sea.dim_out
 
     def compute_input_stats(self, merged: List[dict], path: Optional[DPPath] = None):
         """Update mean and stddev for descriptor elements."""
         return self.sea.compute_input_stats(merged, path)
 
     @classmethod
     def get_data_process_key(cls, config):
@@ -313,6 +316,7 @@
                     resnet_dt=self.resnet_dt,
                 )
             self.filter_layers = filter_layers
+        self.stats = None
 
     def get_rcut(self) -> float:
         """Returns the cut-off radius."""
@@ -374,20 +378,29 @@
         else:
             raise KeyError(key)
 
     def compute_input_stats(self, merged: List[dict], path: Optional[DPPath] = None):
         """Update mean and stddev for descriptor elements."""
         env_mat_stat = EnvMatStatSeA(self)
         if path is not None:
             path = path / env_mat_stat.get_hash()
         env_mat_stat.load_or_compute_stats(merged, path)
+        self.stats = env_mat_stat.stats
         mean, stddev = env_mat_stat()
         if not self.set_davg_zero:
             self.mean.copy_(torch.tensor(mean, device=env.DEVICE))
         self.stddev.copy_(torch.tensor(stddev, device=env.DEVICE))
         if not self.set_davg_zero:
             self.mean.copy_(torch.tensor(mean, device=env.DEVICE))
         self.stddev.copy_(torch.tensor(stddev, device=env.DEVICE))
 
+    def get_stats(self) -> dict[str, StatItem]:
+        """Get the statistics of the descriptor."""
+        if self.stats is None:
+            raise RuntimeError(
+                "The statistics of the descriptor has not been computed."
+            )
+        return self.stats
+
     def forward(
         self,
         nlist: torch.Tensor,

deepmd/pt/model/descriptor/se_atten.py

@@ -20,10 +20,13 @@
 from deepmd.pt.utils import (
     env,
 )
 from deepmd.pt.utils.env_mat_stat import (
     EnvMatStatSeA,
 )
+from deepmd.utils.env_mat_stat import (
+    StatItem,
+)
 from deepmd.utils.path import (
     DPPath,
 )
 
@@ -137,6 +140,7 @@
         )
         filter_layers.append(one)
         self.filter_layers = torch.nn.ModuleList(filter_layers)
+        self.stats = None
 
     def get_rcut(self) -> float:
         """Returns the cut-off radius."""
@@ -187,17 +191,26 @@
         """Returns the output dimension of embedding."""
         return self.get_dim_emb()
 
     def compute_input_stats(self, merged: List[dict], path: Optional[DPPath] = None):
         """Update mean and stddev for descriptor elements."""
         env_mat_stat = EnvMatStatSeA(self)
         if path is not None:
             path = path / env_mat_stat.get_hash()
         env_mat_stat.load_or_compute_stats(merged, path)
+        self.stats = env_mat_stat.stats
         mean, stddev = env_mat_stat()
         if not self.set_davg_zero:
             self.mean.copy_(torch.tensor(mean, device=env.DEVICE))
         self.stddev.copy_(torch.tensor(stddev, device=env.DEVICE))
 
+    def get_stats(self) -> dict[str, StatItem]:
+        """Get the statistics of the descriptor."""
+        if self.stats is None:
+            raise RuntimeError(
+                "The statistics of the descriptor has not been computed."
+            )
+        return self.stats
+
     def forward(
         self,
         nlist: torch.Tensor,