From 9c9f4631ec7e3a9bf7c5d733703ff89a780cf150 Mon Sep 17 00:00:00 2001 From: Jia-Xin Zhu Date: Sun, 4 Jan 2026 00:23:18 +0800 Subject: [PATCH 1/6] optimize tf modifier in deepeval --- deepmd/tf/infer/deep_eval.py | 41 +++++++---------------------- deepmd/tf/modifier/dipole_charge.py | 35 ++++++++++++++++++++++++ 2 files changed, 45 insertions(+), 31 deletions(-) diff --git a/deepmd/tf/infer/deep_eval.py b/deepmd/tf/infer/deep_eval.py index d1eca83055..fc010eb509 100644 --- a/deepmd/tf/infer/deep_eval.py +++ b/deepmd/tf/infer/deep_eval.py @@ -137,37 +137,15 @@ def __init__( self.has_aparam = self.tensors["aparam"] is not None self.has_spin = self.ntypes_spin > 0 - # looks ugly... - if self.modifier_type == "dipole_charge": - from deepmd.tf.modifier import ( - DipoleChargeModifier, - ) + from deepmd.tf.modifier import ( + BaseModifier, + ) - t_mdl_name = self._get_tensor("modifier_attr/mdl_name:0") - t_mdl_charge_map = self._get_tensor("modifier_attr/mdl_charge_map:0") - t_sys_charge_map = self._get_tensor("modifier_attr/sys_charge_map:0") - t_ewald_h = self._get_tensor("modifier_attr/ewald_h:0") - t_ewald_beta = self._get_tensor("modifier_attr/ewald_beta:0") - [mdl_name, mdl_charge_map, sys_charge_map, ewald_h, ewald_beta] = run_sess( - self.sess, - [ - t_mdl_name, - t_mdl_charge_map, - t_sys_charge_map, - t_ewald_h, - t_ewald_beta, - ], - ) - mdl_name = mdl_name.decode("UTF-8") - mdl_charge_map = [int(ii) for ii in mdl_charge_map.decode("UTF-8").split()] - sys_charge_map = [int(ii) for ii in sys_charge_map.decode("UTF-8").split()] - self.dm = DipoleChargeModifier( - mdl_name, - mdl_charge_map, - sys_charge_map, - ewald_h=ewald_h, - ewald_beta=ewald_beta, - ) + self.dm = None + if self.modifier_type is not None: + modifier = BaseModifier.get_class_by_type(self.modifier_type) + modifier_params = modifier.get_params(self) + self.dm = modifier.get_modifier(modifier_params) def _init_tensors(self) -> None: tensor_names = { @@ -684,7 +662,8 @@ def _get_natoms_and_nframes( coords: np.ndarray, atom_types: list[int] | np.ndarray, ) -> tuple[int, int]: - natoms = len(atom_types[0]) + atom_types = np.reshape(atom_types, (-1)) + natoms = len(atom_types) if natoms == 0: assert coords.size == 0 else: diff --git a/deepmd/tf/modifier/dipole_charge.py b/deepmd/tf/modifier/dipole_charge.py index d40c9ccd2f..7a9adf7cf3 100644 --- a/deepmd/tf/modifier/dipole_charge.py +++ b/deepmd/tf/modifier/dipole_charge.py @@ -13,6 +13,9 @@ op_module, tf, ) +from deepmd.tf.infer import ( + DeepEval, +) from deepmd.tf.infer.deep_dipole import DeepDipoleOld as DeepDipole from deepmd.tf.infer.ewald_recp import ( EwaldRecp, @@ -487,3 +490,35 @@ def modify_data(self, data: dict, data_sys: DeepmdData) -> None: data["force"] -= tot_f.reshape(data["force"].shape) if "find_virial" in data and data["find_virial"] == 1.0: data["virial"] -= tot_v.reshape(data["virial"].shape) + + @staticmethod + def get_params(model: DeepEval): + t_mdl_name = model._get_tensor("modifier_attr/mdl_name:0") + t_mdl_charge_map = model._get_tensor("modifier_attr/mdl_charge_map:0") + t_sys_charge_map = model._get_tensor("modifier_attr/sys_charge_map:0") + t_ewald_h = model._get_tensor("modifier_attr/ewald_h:0") + t_ewald_beta = model._get_tensor("modifier_attr/ewald_beta:0") + [mdl_name, mdl_charge_map, sys_charge_map, ewald_h, ewald_beta] = run_sess( + model.sess, + [ + t_mdl_name, + t_mdl_charge_map, + t_sys_charge_map, + t_ewald_h, + t_ewald_beta, + ], + ) + model_charge_map = [ + int(float(ii)) for ii in mdl_charge_map.decode("UTF-8").split() + ] + sys_charge_map = [ + int(float(ii)) for ii in sys_charge_map.decode("UTF-8").split() + ] + modifier_params = { + "model_name": mdl_name.decode("UTF-8"), + "model_charge_map": model_charge_map, + "sys_charge_map": sys_charge_map, + "ewald_h": ewald_h, + "ewald_beta": ewald_beta, + } + return modifier_params From e58b5fc4cd646df33fdc342ca41a4e0110498f8e Mon Sep 17 00:00:00 2001 From: Jia-Xin Zhu Date: Sun, 4 Jan 2026 10:07:57 +0800 Subject: [PATCH 2/6] Fix bug and clean up codes based on coderabbit --- deepmd/tf/infer/deep_eval.py | 6 +++--- deepmd/tf/modifier/base_modifier.py | 24 +++++++++++++++++++++++ deepmd/tf/modifier/dipole_charge.py | 30 +++++++++++++++++++++++++---- 3 files changed, 53 insertions(+), 7 deletions(-) diff --git a/deepmd/tf/infer/deep_eval.py b/deepmd/tf/infer/deep_eval.py index fc010eb509..b50cbb6ed4 100644 --- a/deepmd/tf/infer/deep_eval.py +++ b/deepmd/tf/infer/deep_eval.py @@ -144,7 +144,7 @@ def __init__( self.dm = None if self.modifier_type is not None: modifier = BaseModifier.get_class_by_type(self.modifier_type) - modifier_params = modifier.get_params(self) + modifier_params = modifier.get_params_from_frozen_model(self) self.dm = modifier.get_modifier(modifier_params) def _init_tensors(self) -> None: @@ -662,8 +662,8 @@ def _get_natoms_and_nframes( coords: np.ndarray, atom_types: list[int] | np.ndarray, ) -> tuple[int, int]: - atom_types = np.reshape(atom_types, (-1)) - natoms = len(atom_types) + # (natoms,) or (nframes, natoms,) + natoms = np.shape(atom_types)[-1] if natoms == 0: assert coords.size == 0 else: diff --git a/deepmd/tf/modifier/base_modifier.py b/deepmd/tf/modifier/base_modifier.py index 4e214e0835..167811cd2a 100644 --- a/deepmd/tf/modifier/base_modifier.py +++ b/deepmd/tf/modifier/base_modifier.py @@ -1,4 +1,8 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from abc import ( + abstractmethod, +) + from deepmd.dpmodel.modifier.base_modifier import ( make_base_modifier, ) @@ -11,3 +15,23 @@ class BaseModifier(DeepPot, make_base_modifier()): def __init__(self, *args, **kwargs) -> None: """Construct a basic model for different tasks.""" DeepPot.__init__(self, *args, **kwargs) + + @staticmethod + @abstractmethod + def get_params_from_frozen_model(model) -> dict: + """Extract the modifier parameters from a model. + + This method should extract the necessary parameters from a model + to create an instance of this modifier. + + Parameters + ---------- + model + The model from which to extract parameters + + Returns + ------- + dict + The modifier parameters + """ + pass diff --git a/deepmd/tf/modifier/dipole_charge.py b/deepmd/tf/modifier/dipole_charge.py index 7a9adf7cf3..a2d3efb353 100644 --- a/deepmd/tf/modifier/dipole_charge.py +++ b/deepmd/tf/modifier/dipole_charge.py @@ -1,5 +1,8 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import os +from typing import ( + TYPE_CHECKING, +) import numpy as np @@ -13,9 +16,6 @@ op_module, tf, ) -from deepmd.tf.infer import ( - DeepEval, -) from deepmd.tf.infer.deep_dipole import DeepDipoleOld as DeepDipole from deepmd.tf.infer.ewald_recp import ( EwaldRecp, @@ -30,6 +30,11 @@ run_sess, ) +if TYPE_CHECKING: + from deepmd.tf.infer import ( + DeepEval, + ) + @BaseModifier.register("dipole_charge") class DipoleChargeModifier(DeepDipole, BaseModifier): @@ -492,7 +497,24 @@ def modify_data(self, data: dict, data_sys: DeepmdData) -> None: data["virial"] -= tot_v.reshape(data["virial"].shape) @staticmethod - def get_params(model: DeepEval): + def get_params_from_frozen_model(model: "DeepEval") -> dict: + """Extract modifier parameters from a DeepEval model. + + Parameters + ---------- + model : DeepEval + The DeepEval model instance containing the modifier tensors. + + Returns + ------- + dict + Dictionary containing modifier parameters: + - model_name : str + - model_charge_map : list[int] + - sys_charge_map : list[int] + - ewald_h : float + - ewald_beta : float + """ t_mdl_name = model._get_tensor("modifier_attr/mdl_name:0") t_mdl_charge_map = model._get_tensor("modifier_attr/mdl_charge_map:0") t_sys_charge_map = model._get_tensor("modifier_attr/sys_charge_map:0") From 3069b32f65391db8d96788cf08de0e1b83864130 Mon Sep 17 00:00:00 2001 From: Jia-Xin Zhu Date: Sun, 4 Jan 2026 12:09:55 +0800 Subject: [PATCH 3/6] add exception capture for modifier loading --- deepmd/tf/infer/deep_eval.py | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/deepmd/tf/infer/deep_eval.py b/deepmd/tf/infer/deep_eval.py index b50cbb6ed4..91c0d0b648 100644 --- a/deepmd/tf/infer/deep_eval.py +++ b/deepmd/tf/infer/deep_eval.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +import logging import json from collections.abc import ( Callable, @@ -51,6 +52,8 @@ run_sess, ) +log = logging.getLogger(__name__) + if TYPE_CHECKING: from pathlib import ( Path, @@ -143,9 +146,17 @@ def __init__( self.dm = None if self.modifier_type is not None: - modifier = BaseModifier.get_class_by_type(self.modifier_type) - modifier_params = modifier.get_params_from_frozen_model(self) - self.dm = modifier.get_modifier(modifier_params) + try: + modifier = BaseModifier.get_class_by_type(self.modifier_type) + modifier_params = modifier.get_params_from_frozen_model(self) + self.dm = modifier.get_modifier(modifier_params) + except Exception as exc: + log.warning( + f"Failed to load data modifier '{self.modifier_type}'. " + "The model will be loaded without the data modifier. " + f"Error details: {exc}" + ) + self.modifier_type = None def _init_tensors(self) -> None: tensor_names = { From 4038aed9f688aa733049e7b452c1bacce3d8e8f7 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 4 Jan 2026 04:11:41 +0000 Subject: [PATCH 4/6] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/tf/infer/deep_eval.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepmd/tf/infer/deep_eval.py b/deepmd/tf/infer/deep_eval.py index 91c0d0b648..26d108774b 100644 --- a/deepmd/tf/infer/deep_eval.py +++ b/deepmd/tf/infer/deep_eval.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -import logging import json +import logging from collections.abc import ( Callable, ) From 54f47d3378849b348368ffbdbd8ba988ef1fadae Mon Sep 17 00:00:00 2001 From: Jia-Xin Zhu Date: Mon, 5 Jan 2026 16:21:06 +0800 Subject: [PATCH 5/6] add `skip_modifier` for tf model inference If skip_modifier = True, the data modifier will not be loaded. default value is set as False. --- deepmd/tf/infer/deep_eval.py | 17 ++++++----------- 1 file changed, 6 insertions(+), 11 deletions(-) diff --git a/deepmd/tf/infer/deep_eval.py b/deepmd/tf/infer/deep_eval.py index 26d108774b..c54ecde695 100644 --- a/deepmd/tf/infer/deep_eval.py +++ b/deepmd/tf/infer/deep_eval.py @@ -140,23 +140,18 @@ def __init__( self.has_aparam = self.tensors["aparam"] is not None self.has_spin = self.ntypes_spin > 0 + if kwargs.get("skip_modifier", False): + self.modifier_type = None + from deepmd.tf.modifier import ( BaseModifier, ) self.dm = None if self.modifier_type is not None: - try: - modifier = BaseModifier.get_class_by_type(self.modifier_type) - modifier_params = modifier.get_params_from_frozen_model(self) - self.dm = modifier.get_modifier(modifier_params) - except Exception as exc: - log.warning( - f"Failed to load data modifier '{self.modifier_type}'. " - "The model will be loaded without the data modifier. " - f"Error details: {exc}" - ) - self.modifier_type = None + modifier = BaseModifier.get_class_by_type(self.modifier_type) + modifier_params = modifier.get_params_from_frozen_model(self) + self.dm = modifier.get_modifier(modifier_params) def _init_tensors(self) -> None: tensor_names = { From 90360840f071c87465d5691cb04ecc57b529ad9b Mon Sep 17 00:00:00 2001 From: Jia-Xin Zhu Date: Mon, 5 Jan 2026 16:28:54 +0800 Subject: [PATCH 6/6] minor optimization based on @coderabbit --- deepmd/tf/infer/deep_eval.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/deepmd/tf/infer/deep_eval.py b/deepmd/tf/infer/deep_eval.py index c54ecde695..20f513cd75 100644 --- a/deepmd/tf/infer/deep_eval.py +++ b/deepmd/tf/infer/deep_eval.py @@ -149,9 +149,16 @@ def __init__( self.dm = None if self.modifier_type is not None: - modifier = BaseModifier.get_class_by_type(self.modifier_type) - modifier_params = modifier.get_params_from_frozen_model(self) - self.dm = modifier.get_modifier(modifier_params) + try: + modifier = BaseModifier.get_class_by_type(self.modifier_type) + modifier_params = modifier.get_params_from_frozen_model(self) + self.dm = modifier.get_modifier(modifier_params) + except Exception as exc: + raise RuntimeError( + f"Failed to load data modifier '{self.modifier_type}'. " + f"Use skip_modifier=True to load the model without the modifier. " + f"Error: {exc}" + ) from exc def _init_tensors(self) -> None: tensor_names = {