diff --git a/deepmd/infer/deep_property.py b/deepmd/infer/deep_property.py index 5e35dcd781..19e613191e 100644 --- a/deepmd/infer/deep_property.py +++ b/deepmd/infer/deep_property.py @@ -141,6 +141,18 @@ def eval( nframes, self.get_task_dim() ) + # --- softmax-weighted averaging over frames (minimal) --- + print(f"Nframes == {nframes}") + if nframes != 3: + raise RuntimeError(f"Expected nframes == 3, got {nframes}") + scores = property.mean(axis=1) # (3,) + # If you want to favor *smaller* values (e.g., energies), use: scores = -scores + w = np.exp(scores - scores.max()) + w /= w.sum() # (3,) + avg = (w[:, None] * property).sum(axis=0, keepdims=True) # (1, D) + property[:] = np.repeat(avg, nframes, axis=0) # (3, D) + # -------------------------------------------------------- + if atomic: return ( property, diff --git a/deepmd/pt/loss/property.py b/deepmd/pt/loss/property.py index 189bcb2a4a..c09cdee3ed 100644 --- a/deepmd/pt/loss/property.py +++ b/deepmd/pt/loss/property.py @@ -99,6 +99,27 @@ def forward( """ model_pred = model(**input_dict) var_name = self.var_name + + # ---- Softmax-weighted averaging over the batch added by YL---- + # model_pred[var_name]: (nbz, task_dim) + # 1) get a scalar score per sample (mean over task_dim) + # (If you want to favor smaller values, use `score_per_sample = -model_pred[var_name].mean(dim=1)`.) + score_per_sample = model_pred[var_name].mean(dim=1) # (nbz,) + weights = F.softmax(score_per_sample, dim=0) # (nbz,) + # 2) weighted average vector (1, task_dim) + avg_vec = (weights.unsqueeze(1) * model_pred[var_name]).sum(dim=0, keepdim=True) + # 3) replace all predictions with the averaged vector (broadcast over batch) + model_pred[var_name] = avg_vec.expand_as(model_pred[var_name]) + # ---------------------------------------------------- + + nbz = model_pred[var_name].shape[0] + # =======Raise error when nbz!=3======= + if nbz != 3: + raise RuntimeError( + f"[PropertyLoss] Expected batch size nbz == 3 for softmax-avg, got nbz == {nbz}. " + "Ensure your DataLoader yields triples (batch_size=3, drop_last=True)." + ) + nbz = model_pred[var_name].shape[0] assert model_pred[var_name].shape == (nbz, self.task_dim) assert label[var_name].shape == (nbz, self.task_dim) diff --git a/deepmd/utils/path.py b/deepmd/utils/path.py index b9e48fe531..fd007632e1 100644 --- a/deepmd/utils/path.py +++ b/deepmd/utils/path.py @@ -331,7 +331,7 @@ def _load_h5py(cls, path: str, mode: str = "r") -> h5py.File: # this method has cache to avoid duplicated # loading from different DPH5Path # However the file will be never closed? - return h5py.File(path, mode) + return h5py.File(path, mode, locking=False) def load_numpy(self) -> np.ndarray: """Load NumPy array.