merge with devel

wanghan-iapcm · Feb 28, 2024 · cd16e7b · cd16e7b
2 parents b48c5ee + 2a1508d
commit cd16e7b
Show file tree

Hide file tree

Showing 64 changed files with 2,996 additions and 172 deletions.
diff --git a/.github/workflows/test_cc.yml b/.github/workflows/test_cc.yml
@@ -56,7 +56,7 @@ jobs:
         TF_INTRA_OP_PARALLELISM_THREADS: 1
         TF_INTER_OP_PARALLELISM_THREADS: 1
         LAMMPS_PLUGIN_PATH: ${{ github.workspace }}/dp_test/lib/deepmd_lmp
-        LD_LIBRARY_PATH: ${{ github.workspace }}/dp_test/lib
+        LD_LIBRARY_PATH: ${{ github.workspace }}/dp_test/lib:${{ github.workspace }}/libtorch/lib
       if: ${{ !matrix.check_memleak }}
     # test ipi
     - run: pytest --cov=deepmd source/ipi/tests
@@ -65,7 +65,7 @@ jobs:
         TF_INTRA_OP_PARALLELISM_THREADS: 1
         TF_INTER_OP_PARALLELISM_THREADS: 1
         PATH: ${{ github.workspace }}/dp_test/bin:$PATH
-        LD_LIBRARY_PATH: ${{ github.workspace }}/dp_test/lib
+        LD_LIBRARY_PATH: ${{ github.workspace }}/dp_test/lib:${{ github.workspace }}/libtorch/lib
       if: ${{ !matrix.check_memleak }}
     - uses: codecov/codecov-action@v4
       env:

diff --git a/.github/workflows/test_cuda.yml b/.github/workflows/test_cuda.yml
@@ -38,6 +38,8 @@ jobs:
       with:
         useLocalCache: true
         useCloudCache: false
+    - name: Install wget and unzip
+      run: apt-get update && apt-get install -y wget unzip
     - run: |
          wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.0-1_all.deb \
          && sudo dpkg -i cuda-keyring_1.0-1_all.deb \
@@ -53,7 +55,13 @@ jobs:
         DP_ENABLE_NATIVE_OPTIMIZATION: 1
     - run: dp --version
     - run: python -m pytest source/tests --durations=0
-    - run: source/install/test_cc_local.sh
+    - name: Download libtorch
+      run: |
+         wget https://download.pytorch.org/libtorch/cu121/libtorch-cxx11-abi-shared-with-deps-2.2.1%2Bcu121.zip -O libtorch.zip
+         unzip libtorch.zip
+    - run: |
+        export CMAKE_PREFIX_PATH=$GITHUB_WORKSPACE/libtorch
+        source/install/test_cc_local.sh
       env:
         OMP_NUM_THREADS: 1
         TF_INTRA_OP_PARALLELISM_THREADS: 1
@@ -63,7 +71,7 @@ jobs:
         DP_VARIANT: cuda
         DP_USE_MPICH2: 1
     - run: |
-        export LD_LIBRARY_PATH=$GITHUB_WORKSPACE/dp_test/lib:$CUDA_PATH/lib64:$LD_LIBRARY_PATH
+        export LD_LIBRARY_PATH=$GITHUB_WORKSPACE/dp_test/lib:$GITHUB_WORKSPACE/libtorch/lib:$CUDA_PATH/lib64:$LD_LIBRARY_PATH
         export PATH=$GITHUB_WORKSPACE/dp_test/bin:$PATH
         python -m pytest source/lmp/tests
         python -m pytest source/ipi/tests

diff --git a/deepmd/dpmodel/atomic_model/dp_atomic_model.py b/deepmd/dpmodel/atomic_model/dp_atomic_model.py
@@ -17,6 +17,9 @@
 from deepmd.dpmodel.output_def import (
     FittingOutputDef,
 )
+from deepmd.utils.version import (
+    check_version_compatibility,
+)
 
 from .base_atomic_model import (
     BaseAtomicModel,
@@ -132,6 +135,7 @@ def serialize(self) -> dict:
         return {
             "@class": "Model",
             "type": "standard",
+            "@version": 1,
             "type_map": self.type_map,
             "descriptor": self.descriptor.serialize(),
             "fitting": self.fitting.serialize(),
@@ -140,6 +144,7 @@ def serialize(self) -> dict:
     @classmethod
     def deserialize(cls, data) -> "DPAtomicModel":
         data = copy.deepcopy(data)
+        check_version_compatibility(data.pop("@version", 1), 1, 1)
         data.pop("@class")
         data.pop("type")
         descriptor_obj = BaseDescriptor.deserialize(data["descriptor"])

diff --git a/deepmd/dpmodel/atomic_model/linear_atomic_model.py b/deepmd/dpmodel/atomic_model/linear_atomic_model.py
@@ -19,6 +19,9 @@
     get_multiple_nlist_key,
     nlist_distinguish_types,
 )
+from deepmd.utils.version import (
+    check_version_compatibility,
+)
 
 from ..output_def import (
     FittingOutputDef,
@@ -185,13 +188,15 @@ def serialize(models) -> dict:
         return {
             "@class": "Model",
             "type": "linear",
+            "@version": 1,
             "models": [model.serialize() for model in models],
             "model_name": [model.__class__.__name__ for model in models],
         }
 
     @staticmethod
     def deserialize(data) -> List[BaseAtomicModel]:
         data = copy.deepcopy(data)
+        check_version_compatibility(data.pop("@version", 1), 1, 1)
         data.pop("@class")
         data.pop("type")
         model_names = data["model_name"]
@@ -271,6 +276,7 @@ def serialize(self) -> dict:
         return {
             "@class": "Model",
             "type": "zbl",
+            "@version": 1,
             "models": LinearAtomicModel.serialize([self.dp_model, self.zbl_model]),
             "sw_rmin": self.sw_rmin,
             "sw_rmax": self.sw_rmax,
@@ -280,6 +286,7 @@ def serialize(self) -> dict:
     @classmethod
     def deserialize(cls, data) -> "DPZBLLinearAtomicModel":
         data = copy.deepcopy(data)
+        check_version_compatibility(data.pop("@version", 1), 1, 1)
         data.pop("@class")
         data.pop("type")
         sw_rmin = data["sw_rmin"]

diff --git a/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py b/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py
@@ -16,6 +16,9 @@
 from deepmd.utils.pair_tab import (
     PairTab,
 )
+from deepmd.utils.version import (
+    check_version_compatibility,
+)
 
 from .base_atomic_model import (
     BaseAtomicModel,
@@ -109,6 +112,7 @@ def serialize(self) -> dict:
         return {
             "@class": "Model",
             "type": "pairtab",
+            "@version": 1,
             "tab": self.tab.serialize(),
             "rcut": self.rcut,
             "sel": self.sel,
@@ -117,6 +121,7 @@ def serialize(self) -> dict:
     @classmethod
     def deserialize(cls, data) -> "PairTabAtomicModel":
         data = copy.deepcopy(data)
+        check_version_compatibility(data.pop("@version", 1), 1, 1)
         data.pop("@class")
         data.pop("type")
         rcut = data["rcut"]

diff --git a/deepmd/dpmodel/descriptor/se_e2_a.py b/deepmd/dpmodel/descriptor/se_e2_a.py
@@ -9,6 +9,9 @@
 from deepmd.utils.path import (
     DPPath,
 )
+from deepmd.utils.version import (
+    check_version_compatibility,
+)
 
 try:
     from deepmd._version import version as __version__
@@ -345,6 +348,7 @@ def serialize(self) -> dict:
         return {
             "@class": "Descriptor",
             "type": "se_e2_a",
+            "@version": 1,
             "rcut": self.rcut,
             "rcut_smth": self.rcut_smth,
             "sel": self.sel,
@@ -371,6 +375,7 @@ def serialize(self) -> dict:
     def deserialize(cls, data: dict) -> "DescrptSeA":
         """Deserialize from dict."""
         data = copy.deepcopy(data)
+        check_version_compatibility(data.pop("@version", 1), 1, 1)
         data.pop("@class", None)
         data.pop("type", None)
         variables = data.pop("@variables")

diff --git a/deepmd/dpmodel/descriptor/se_r.py b/deepmd/dpmodel/descriptor/se_r.py
@@ -4,6 +4,9 @@
 from deepmd.utils.path import (
     DPPath,
 )
+from deepmd.utils.version import (
+    check_version_compatibility,
+)
 
 try:
     from deepmd._version import version as __version__
@@ -282,6 +285,7 @@ def serialize(self) -> dict:
         return {
             "@class": "Descriptor",
             "type": "se_r",
+            "@version": 1,
             "rcut": self.rcut,
             "rcut_smth": self.rcut_smth,
             "sel": self.sel,
@@ -307,6 +311,7 @@ def serialize(self) -> dict:
     def deserialize(cls, data: dict) -> "DescrptSeR":
         """Deserialize from dict."""
         data = copy.deepcopy(data)
+        check_version_compatibility(data.pop("@version", 1), 1, 1)
         data.pop("@class", None)
         data.pop("type", None)
         variables = data.pop("@variables")

diff --git a/deepmd/dpmodel/fitting/general_fitting.py b/deepmd/dpmodel/fitting/general_fitting.py
@@ -21,6 +21,9 @@
     FittingNet,
     NetworkCollection,
 )
+from deepmd.utils.version import (
+    check_version_compatibility,
+)
 
 from .base_fitting import (
     BaseFitting,
@@ -210,6 +213,7 @@ def serialize(self) -> dict:
         """Serialize the fitting to dict."""
         return {
             "@class": "Fitting",
+            "@version": 1,
             "var_name": self.var_name,
             "ntypes": self.ntypes,
             "dim_descrpt": self.dim_descrpt,
@@ -241,6 +245,7 @@ def serialize(self) -> dict:
     @classmethod
     def deserialize(cls, data: dict) -> "GeneralFitting":
         data = copy.deepcopy(data)
+        check_version_compatibility(data.pop("@version", 1), 1, 1)
         data.pop("@class")
         data.pop("type")
         variables = data.pop("@variables")

diff --git a/deepmd/dpmodel/utils/network.py b/deepmd/dpmodel/utils/network.py
@@ -20,6 +20,10 @@
 import h5py
 import numpy as np
 
+from deepmd.utils.version import (
+    check_version_compatibility,
+)
+
 try:
     from deepmd._version import version as __version__
 except ImportError:
@@ -189,6 +193,8 @@ def serialize(self) -> dict:
             "idt": self.idt,
         }
         return {
+            "@class": "Layer",
+            "@version": 1,
             "bias": self.b is not None,
             "use_timestep": self.idt is not None,
             "activation_function": self.activation_function,
@@ -208,6 +214,8 @@ def deserialize(cls, data: dict) -> "NativeLayer":
             The dict to deserialize from.
         """
         data = copy.deepcopy(data)
+        check_version_compatibility(data.pop("@version", 1), 1, 1)
+        data.pop("@class", None)
         variables = data.pop("@variables")
         assert variables["w"] is not None and len(variables["w"].shape) == 2
         num_in, num_out = variables["w"].shape
@@ -349,7 +357,11 @@ def serialize(self) -> dict:
             dict
                 The serialized network.
             """
-            return {"layers": [layer.serialize() for layer in self.layers]}
+            return {
+                "@class": "NN",
+                "@version": 1,
+                "layers": [layer.serialize() for layer in self.layers],
+            }
 
         @classmethod
         def deserialize(cls, data: dict) -> "NN":
@@ -360,6 +372,9 @@ def deserialize(cls, data: dict) -> "NN":
             data : dict
                 The dict to deserialize from.
             """
+            data = data.copy()
+            check_version_compatibility(data.pop("@version", 1), 1, 1)
+            data.pop("@class", None)
             return cls(data["layers"])
 
         def __getitem__(self, key):
@@ -471,6 +486,8 @@ def serialize(self) -> dict:
                 The serialized network.
             """
             return {
+                "@class": "EmbeddingNetwork",
+                "@version": 1,
                 "in_dim": self.in_dim,
                 "neuron": self.neuron.copy(),
                 "activation_function": self.activation_function,
@@ -490,6 +507,8 @@ def deserialize(cls, data: dict) -> "EmbeddingNet":
                 The dict to deserialize from.
             """
             data = copy.deepcopy(data)
+            check_version_compatibility(data.pop("@version", 1), 1, 1)
+            data.pop("@class", None)
             layers = data.pop("layers")
             obj = cls(**data)
             super(EN, obj).__init__(layers)
@@ -566,6 +585,8 @@ def serialize(self) -> dict:
                 The serialized network.
             """
             return {
+                "@class": "FittingNetwork",
+                "@version": 1,
                 "in_dim": self.in_dim,
                 "out_dim": self.out_dim,
                 "neuron": self.neuron.copy(),
@@ -586,6 +607,8 @@ def deserialize(cls, data: dict) -> "FittingNet":
                 The dict to deserialize from.
             """
             data = copy.deepcopy(data)
+            check_version_compatibility(data.pop("@version", 1), 1, 1)
+            data.pop("@class", None)
             layers = data.pop("layers")
             obj = cls(**data)
             T_Network.__init__(obj, layers)
@@ -688,6 +711,8 @@ def serialize(self) -> dict:
         network_type_map_inv = {v: k for k, v in self.NETWORK_TYPE_MAP.items()}
         network_type_name = network_type_map_inv[self.network_type]
         return {
+            "@class": "NetworkCollection",
+            "@version": 1,
             "ndim": self.ndim,
             "ntypes": self.ntypes,
             "network_type": network_type_name,
@@ -703,4 +728,7 @@ def deserialize(cls, data: dict) -> "NetworkCollection":
         data : dict
             The dict to deserialize from.
         """
+        data = data.copy()
+        check_version_compatibility(data.pop("@version", 1), 1, 1)
+        data.pop("@class", None)
         return cls(**data)
diff --git a/deepmd/infer/deep_eval.py b/deepmd/infer/deep_eval.py
@@ -472,6 +472,33 @@ def _standard_input(self, coords, cells, atom_types, fparam, aparam, mixed_type)
             aparam = np.array(aparam)
         natoms, nframes = self._get_natoms_and_nframes(coords, atom_types, mixed_type)
         atom_types = self._expande_atype(atom_types, nframes, mixed_type)
+        coords = coords.reshape(nframes, natoms, 3)
+        if cells is not None:
+            cells = cells.reshape(nframes, 3, 3)
+        if fparam is not None:
+            fdim = self.get_dim_fparam()
+            if fparam.size == nframes * fdim:
+                fparam = np.reshape(fparam, [nframes, fdim])
+            elif fparam.size == fdim:
+                fparam = np.tile(fparam.reshape([-1]), [nframes, 1])
+            else:
+                raise RuntimeError(
+                    "got wrong size of frame param, should be either %d x %d or %d"
+                    % (nframes, fdim, fdim)
+                )
+        if aparam is not None:
+            fdim = self.get_dim_aparam()
+            if aparam.size == nframes * natoms * fdim:
+                aparam = np.reshape(aparam, [nframes, natoms * fdim])
+            elif aparam.size == natoms * fdim:
+                aparam = np.tile(aparam.reshape([-1]), [nframes, 1])
+            elif aparam.size == fdim:
+                aparam = np.tile(aparam.reshape([-1]), [nframes, natoms])
+            else:
+                raise RuntimeError(
+                    "got wrong size of frame param, should be either %d x %d x %d or %d x %d or %d"
+                    % (nframes, natoms, fdim, natoms, fdim, fdim)
+                )
         return coords, cells, atom_types, fparam, aparam, nframes, natoms
 
     def get_sel_type(self) -> List[int]:

diff --git a/deepmd/main.py b/deepmd/main.py
@@ -226,7 +226,7 @@ def main_parser() -> argparse.ArgumentParser:
         "--init-frz-model",
         type=str,
         default=None,
-        help="(Supported backend: TensorFlow) Initialize the training from the frozen model.",
+        help="Initialize the training from the frozen model.",
     )
     parser_train_subgroup.add_argument(
         "-t",

diff --git a/deepmd/pt/entrypoints/main.py b/deepmd/pt/entrypoints/main.py
@@ -65,6 +65,7 @@ def get_trainer(
     finetune_model=None,
     model_branch="",
     force_load=False,
+    init_frz_model=None,
 ):
     # Initialize DDP
     local_rank = os.environ.get("LOCAL_RANK")
@@ -200,6 +201,7 @@ def prepare_trainer_input_single(
         finetune_model=finetune_model,
         force_load=force_load,
         shared_links=shared_links,
+        init_frz_model=init_frz_model,
     )
     return trainer
 
@@ -243,6 +245,7 @@ def train(FLAGS):
         FLAGS.finetune,
         FLAGS.model_branch,
         FLAGS.force_load,
+        FLAGS.init_frz_model,
     )
     trainer.run()