Nosample

deepmd_pt/entrypoints/main.py

@@ -64,8 +64,9 @@ def get_trainer(config, ckpt=None, force_load=False, finetune_model=None):
     else:
         train_data = DpLoaderSet(training_systems, training_dataset_params['batch_size'], model_params,
                                  type_split=type_split)
-        data_stat_nbatch = model_params.get('data_stat_nbatch', 10)
-        sampled = make_stat_input(train_data.systems, train_data.dataloaders, data_stat_nbatch)
+        #data_stat_nbatch = model_params.get('data_stat_nbatch', 10)
+        #sampled = make_stat_input(train_data.systems, train_data.dataloaders, data_stat_nbatch)
+        sampled = True
         if noise_settings is not None:
             train_data = DpLoaderSet(training_systems, training_dataset_params['batch_size'], model_params,
                                      type_split=type_split, noise_settings=noise_settings)

deepmd_pt/model/descriptor/se_a.py

@@ -2,6 +2,8 @@
 import torch
 
 from deepmd_pt.utils import env
+from deepmd_pt.utils.stat import sample_system
+from deepmd_pt.utils.dataloader import DpLoaderSet
 from deepmd_pt.model.descriptor import prod_env_mat_se_a, Descriptor, compute_std
 
 try:
@@ -66,15 +68,30 @@ def dim_out(self):
         """
         return self.filter_neuron[-1] * self.axis_neuron
 
-    def compute_input_stats(self, merged):
+    def compute_input_stats(self, nbatch, merged: DpLoaderSet,rcond=1e-3):
         """Update mean and stddev for descriptor elements.
         """
-        sumr = []
-        suma = []
-        sumn = []
-        sumr2 = []
-        suma2 = []
-        for system in merged:  # 逐个 system 的分析
+
+        keys = [
+        "coord",
+        "force",
+        "energy",
+        "atype",
+        "natoms",
+        "mapping",
+        "selected",
+        "selected_loc",
+        "selected_type",
+        "shift",
+        ]
+        natoms = []
+        energy = []
+        sumr = None
+        if merged.systems[0].mixed_type:
+            keys.append("real_natoms_vec")
+
+        for item in merged.dataloaders:  #sample from each system, the intermediate results would not be saved
+            system = sample_system(keys, nbatch, item)
             index = system['mapping'].unsqueeze(-1).expand(-1, -1, 3)
             extended_coord = torch.gather(system['coord'], dim=1, index=index)
             extended_coord = extended_coord - system['shift']
@@ -85,17 +102,28 @@ def compute_input_stats(self, merged):
             )
             sysr, sysr2, sysa, sysa2, sysn = analyze_descrpt(env_mat.detach().cpu().numpy(), self.ndescrpt,
                                                              system['natoms'])
-            sumr.append(sysr)
-            suma.append(sysa)
-            sumn.append(sysn)
-            sumr2.append(sysr2)
-            suma2.append(sysa2)
-        sumr = np.sum(sumr, axis=0)
-        suma = np.sum(suma, axis=0)
-        sumn = np.sum(sumn, axis=0)
-        sumr2 = np.sum(sumr2, axis=0)
-        suma2 = np.sum(suma2, axis=0)
-        return sumr, suma, sumn, sumr2, suma2
+
+            energy.append(system['energy'].mean(dim=0, keepdim=True))
+            if merged.systems[0].mixed_type:
+                natoms.append(system['real_natoms_vec'].double().mean(dim=0, keepdim=True))
+            else:
+                natoms.append(system['natoms'].double().mean(dim=0, keepdim=True))
+            if(sumr is None):
+                sumr = np.add(np.zeros_like(sysr),sysr)
+                suma = np.add(np.zeros_like(sysa),sysa)
+                sumn = np.add(np.zeros_like(sysn),sysn)
+                sumr2 = np.add(np.zeros_like(sysr2),sysr2)
+                suma2 = np.add(np.zeros_like(sysa2),sysa2)
+            else:
+                sumr = np.add(sumr,sysr)
+                suma = np.add(suma,sysa)
+                sumn = np.add(sumn,sysn)
+                sumr2 = np.add(sumr2,sysr2)
+                suma2 = np.add(suma2,sysa2)
+        sys_ener = torch.cat(energy).cpu()
+        sys_tynatom = torch.cat(natoms)[:, 2:].cpu()
+        energy_coef, _, _, _ = np.linalg.lstsq(sys_tynatom, sys_ener, rcond)
+        return sumr, suma, sumn, sumr2, suma2,  energy_coef
 
     def init_desc_stat(self, sumr, suma, sumn, sumr2, suma2):
         all_davg = []

deepmd_pt/model/descriptor/se_atten.py

@@ -1,6 +1,7 @@
 import numpy as np
 import torch
-
+from deepmd_pt.utils.stat import sample_system
+from deepmd_pt.utils.dataloader import DpLoaderSet
 from deepmd_pt.utils import env
 from deepmd_pt.model.descriptor import prod_env_mat_se_a, Descriptor, compute_std
 
@@ -110,16 +111,31 @@ def dim_emb(self):
         """
         return self.filter_neuron[-1]
 
-    def compute_input_stats(self, merged):
+    def compute_input_stats(self, nbatch, merged: DpLoaderSet,rcond=1e-3):
         """Update mean and stddev for descriptor elements.
         """
-        sumr = []
-        suma = []
-        sumn = []
-        sumr2 = []
-        suma2 = []
-        mixed_type = 'real_natoms_vec' in merged[0]
-        for system in merged:  # 逐个 system 的分析
+
+        keys = [
+        "coord",
+        "force",
+        "energy",
+        "atype",
+        "natoms",
+        "mapping",
+        "selected",
+        "selected_loc",
+        "selected_type",
+        "shift",
+        ]
+        natoms = []
+        energy = []
+        sumr = None
+        mixed_type = merged.systems[0].mixed_type
+        if mixed_type:
+            keys.append("real_natoms_vec")
+
+        for item in merged.dataloaders:  #sample from each system, the intermediate results would not be saved
+            system = sample_system(keys, nbatch, item)
             index = system['mapping'].unsqueeze(-1).expand(-1, -1, 3)
             extended_coord = torch.gather(system['coord'], dim=1, index=index)
             extended_coord = extended_coord - system['shift']
@@ -135,17 +151,27 @@ def compute_input_stats(self, merged):
                 sysr, sysr2, sysa, sysa2, sysn = analyze_descrpt(env_mat.detach().cpu().numpy(), self.ndescrpt,
                                                                  system['real_natoms_vec'], mixed_type=mixed_type,
                                                                  real_atype=system['atype'].detach().cpu().numpy())
-            sumr.append(sysr)
-            suma.append(sysa)
-            sumn.append(sysn)
-            sumr2.append(sysr2)
-            suma2.append(sysa2)
-        sumr = np.sum(sumr, axis=0)
-        suma = np.sum(suma, axis=0)
-        sumn = np.sum(sumn, axis=0)
-        sumr2 = np.sum(sumr2, axis=0)
-        suma2 = np.sum(suma2, axis=0)
-        return sumr, suma, sumn, sumr2, suma2
+            energy.append(system['energy'].mean(dim=0, keepdim=True))
+            if merged.systems[0].mixed_type:
+                natoms.append(system['real_natoms_vec'].double().mean(dim=0, keepdim=True))
+            else:
+                natoms.append(system['natoms'].double().mean(dim=0, keepdim=True))
+            if(sumr is None):
+                sumr = np.add(np.zeros_like(sysr),sysr)
+                suma = np.add(np.zeros_like(sysa),sysa)
+                sumn = np.add(np.zeros_like(sysn),sysn)
+                sumr2 = np.add(np.zeros_like(sysr2),sysr2)
+                suma2 = np.add(np.zeros_like(sysa2),sysa2)
+            else:
+                sumr = np.add(sumr,sysr)
+                suma = np.add(suma,sysa)
+                sumn = np.add(sumn,sysn)
+                sumr2 = np.add(sumr2,sysr2)
+                suma2 = np.add(suma2,sysa2)
+        sys_ener = torch.cat(energy).cpu()
+        sys_tynatom = torch.cat(natoms)[:, 2:].cpu()
+        energy_coef, _, _, _ = np.linalg.lstsq(sys_tynatom, sys_ener, rcond)
+        return sumr, suma, sumn, sumr2, suma2,  energy_coef
 
     def init_desc_stat(self, sumr, suma, sumn, sumr2, suma2):
         all_davg = []

deepmd_pt/model/descriptor/se_uni.py

@@ -8,6 +8,8 @@
   List, Optional, Tuple,
 )
 from deepmd_pt.utils import env
+from deepmd_pt.utils.stat import sample_system
+from deepmd_pt.utils.dataloader import DpLoaderSet
 from deepmd_pt.utils.utils import get_activation_fn, ActivationFn
 from deepmd_pt.model.descriptor import prod_env_mat_se_a, Descriptor, compute_std
 from deepmd_pt.model.network import (
@@ -529,17 +531,30 @@ def list_update(update_list):
     return g1_new, g2_new, h2_new
 
 
-  def compute_input_stats(self, merged):
+  def compute_input_stats(self, nbatch, merged: DpLoaderSet,rcond=1e-3):
       """Update mean and stddev for descriptor elements.
       """
       ndescrpt = self.nnei * 4  # 描述符的元素数量
-      sumr = []
-      suma = []
-      sumn = []
-      sumr2 = []
-      suma2 = []
-      mixed_type = 'real_natoms_vec' in merged[0]
-      for system in merged:  # 逐个 system 的分析
+      keys = [
+        "coord",
+        "force",
+        "energy",
+        "atype",
+        "natoms",
+        "mapping",
+        "selected",
+        "selected_loc",
+        "selected_type",
+        "shift",
+        ]
+      natoms = []
+      energy = []
+      sumr = None
+      mixed_type = merged.systems[0].mixed_type
+      if mixed_type:
+        keys.append("real_natoms_vec")
+      for item in merged.dataloaders:  #sample from each system, the intermediate results would not be saved
+          system = sample_system(keys, nbatch, item)
           index = system['mapping'].unsqueeze(-1).expand(-1, -1, 3)
           extended_coord = torch.gather(system['coord'], dim=1, index=index)
           extended_coord = extended_coord - system['shift']
@@ -557,17 +572,27 @@ def compute_input_stats(self, merged):
                 env_mat.detach().cpu().numpy(), ndescrpt,
                 system['real_natoms_vec'], mixed_type=mixed_type,
                 real_atype=system['atype'].detach().cpu().numpy())
-          sumr.append(sysr)
-          suma.append(sysa)
-          sumn.append(sysn)
-          sumr2.append(sysr2)
-          suma2.append(sysa2)
-      sumr = np.sum(sumr, axis=0)
-      suma = np.sum(suma, axis=0)
-      sumn = np.sum(sumn, axis=0)
-      sumr2 = np.sum(sumr2, axis=0)
-      suma2 = np.sum(suma2, axis=0)
-      return sumr, suma, sumn, sumr2, suma2
+          energy.append(system['energy'].mean(dim=0, keepdim=True))
+          if merged.systems[0].mixed_type:
+                natoms.append(system['real_natoms_vec'].double().mean(dim=0, keepdim=True))
+          else:
+                natoms.append(system['natoms'].double().mean(dim=0, keepdim=True))
+          if(sumr is None):
+                sumr = np.add(np.zeros_like(sysr),sysr)
+                suma = np.add(np.zeros_like(sysa),sysa)
+                sumn = np.add(np.zeros_like(sysn),sysn)
+                sumr2 = np.add(np.zeros_like(sysr2),sysr2)
+                suma2 = np.add(np.zeros_like(sysa2),sysa2)
+          else:
+                sumr = np.add(sumr,sysr)
+                suma = np.add(suma,sysa)
+                sumn = np.add(sumn,sysn)
+                sumr2 = np.add(sumr2,sysr2)
+                suma2 = np.add(suma2,sysa2)
+      sys_ener = torch.cat(energy).cpu()
+      sys_tynatom = torch.cat(natoms)[:, 2:].cpu()
+      energy_coef, _, _, _ = np.linalg.lstsq(sys_tynatom, sys_ener, rcond)
+      return sumr, suma, sumn, sumr2, suma2,  energy_coef
 
 
   def init_desc_stat(self, sumr, suma, sumn, sumr2, suma2):

deepmd_pt/model/model/__init__.py

@@ -9,27 +9,27 @@
 from .dpau_ener import EnergyModelDPAUni
 
 
-def get_model(model_params, sampled=None):
+def get_model(model_params, training_data=None, sampled=None):
     if model_params.get("fitting_net", None) is not None:
         if model_params.get("backbone", None) is None:
             if model_params["descriptor"]["type"] == "se_e2_a":
-                return EnergyModelSeA(model_params, sampled)
+                return EnergyModelSeA(model_params, training_data, sampled)
             elif model_params["descriptor"]["type"] == "se_atten":
                 if model_params["fitting_net"].get("type", "ener") == "ener":
-                    return EnergyModelDPA1(model_params, sampled)
+                    return EnergyModelDPA1(model_params,training_data, sampled)
                 elif "direct" in model_params["fitting_net"].get("type", "ener"):
-                    return ForceModelDPA1(model_params, sampled)
+                    return ForceModelDPA1(model_params, training_data, sampled)
             elif model_params["descriptor"]["type"] == "se_uni":
-                return EnergyModelDPAUni(model_params, sampled)
+                return EnergyModelDPAUni(model_params, training_data, sampled)
             else:
                 raise NotImplementedError
         else:
             if model_params["fitting_net"].get("type", "ener") == "ener":
-                return EnergyModelDPA2(model_params, sampled)
+                return EnergyModelDPA2(model_params, training_data, sampled)
             elif "direct" in model_params["fitting_net"].get("type", "ener"):
-                return ForceModelDPA2(model_params, sampled)
+                return ForceModelDPA2(model_params, training_data, sampled)
     else:
         if model_params.get("backbone", None) is None:
-            return DenoiseModelDPA1(model_params, sampled)
+            return DenoiseModelDPA1(model_params,training_data, sampled)
         else:
-            return DenoiseModelDPA2(model_params, sampled)
+            return DenoiseModelDPA2(model_params, training_data,sampled)

deepmd_pt/model/model/dpa1_denoise.py

@@ -10,7 +10,7 @@
 
 class DenoiseModelDPA1(BaseModel):
 
-    def __init__(self, model_params, sampled=None):
+    def __init__(self, model_params, training_data, sampled=None):
         """Based on components, construct a DPA-1 model for energy.
 
         Args:
@@ -43,7 +43,7 @@ def __init__(self, model_params, sampled=None):
         self.descriptor = DescrptSeAtten(**descriptor_param)
 
         # Statistics
-        self.compute_or_load_stat(model_params, {}, ntypes, sampled=sampled)
+        self.compute_or_load_stat(model_params, {}, ntypes, training_data, sampled=sampled)
 
         assert model_params.pop('fitting_net', None) is None, f'Denoise task must not have fitting_net!'
         # Denoise and predict

deepmd_pt/model/model/dpa1_ener.py

@@ -14,7 +14,7 @@
 
 class EnergyModelDPA1(BaseModel):
 
-    def __init__(self, model_params, sampled=None):
+    def __init__(self, model_params, training_data, sampled=None):
         """Based on components, construct a DPA-1 model for energy.
 
         Args:
@@ -55,7 +55,7 @@ def __init__(self, model_params, sampled=None):
         fitting_param['use_tebd'] = True
 
         # Statistics
-        self.compute_or_load_stat(model_params, fitting_param, ntypes, sampled=sampled)
+        self.compute_or_load_stat(model_params, fitting_param, ntypes, training_data, sampled=sampled)
 
         self.fitting_net = EnergyFittingNetType(**fitting_param)
 

deepmd_pt/model/model/dpa1_force.py

@@ -13,7 +13,7 @@
 
 class ForceModelDPA1(BaseModel):
 
-    def __init__(self, model_params, sampled=None):
+    def __init__(self, model_params, training_data, sampled=None):
         """Based on components, construct a DPA-1 model for energy.
 
         Args:
@@ -57,7 +57,7 @@ def __init__(self, model_params, sampled=None):
         self.fitting_net_force = DipoleFittingNetType(**fitting_param)
 
         # Statistics
-        self.compute_or_load_stat(model_params, fitting_param, ntypes, sampled=sampled)
+        self.compute_or_load_stat(model_params, fitting_param, ntypes, training_data, sampled=sampled)
 
         if fitting_type == 'direct_force_ener':
             self.fitting_net_ener = EnergyFittingNetType(**fitting_param)

deepmd_pt/model/model/dpa2_denoise.py

@@ -12,7 +12,7 @@
 
 class DenoiseModelDPA2(BaseModel):
 
-    def __init__(self, model_params, sampled=None):
+    def __init__(self, model_params, training_data, sampled=None):
         """Based on components, construct a DPA-1 model for energy.
 
         Args:
@@ -45,7 +45,7 @@ def __init__(self, model_params, sampled=None):
         self.descriptor = DescrptSeAtten(**descriptor_param)
 
         # Statistics
-        self.compute_or_load_stat(model_params, {}, ntypes, sampled=sampled)
+        self.compute_or_load_stat(model_params, {}, ntypes, training_data, sampled=sampled)
 
         # BackBone
         backbone_param = model_params.pop('backbone')