chore: use pickle instead of copy.deepcopy

Signed-off-by: Jinzhe Zeng <[email protected]>

deepmd/dpmodel/atomic_model/dp_atomic_model.py

@@ -1,5 +1,4 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-import copy
 from typing import (
     Optional,
 )
@@ -15,6 +14,9 @@
 from deepmd.dpmodel.output_def import (
     FittingOutputDef,
 )
+from deepmd.utils.copy import (
+    deepcopy,
+)
 from deepmd.utils.version import (
     check_version_compatibility,
 )
@@ -208,7 +210,7 @@ def serialize(self) -> dict:
 
     @classmethod
     def deserialize(cls, data) -> "DPAtomicModel":
-        data = copy.deepcopy(data)
+        data = deepcopy(data)
         check_version_compatibility(data.pop("@version", 1), 2, 2)
         data.pop("@class")
         data.pop("type")

deepmd/dpmodel/atomic_model/linear_atomic_model.py

@@ -1,5 +1,4 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-import copy
 from typing import (
     Optional,
     Union,
@@ -16,6 +15,9 @@
 from deepmd.env import (
     GLOBAL_NP_FLOAT_PRECISION,
 )
+from deepmd.utils.copy import (
+    deepcopy,
+)
 from deepmd.utils.version import (
     check_version_compatibility,
 )
@@ -305,7 +307,7 @@ def serialize(self) -> dict:
 
     @classmethod
     def deserialize(cls, data: dict) -> "LinearEnergyAtomicModel":
-        data = copy.deepcopy(data)
+        data = deepcopy(data)
         check_version_compatibility(data.pop("@version", 2), 2, 2)
         data.pop("@class", None)
         data.pop("type", None)
@@ -418,7 +420,7 @@ def serialize(self) -> dict:
 
     @classmethod
     def deserialize(cls, data) -> "DPZBLLinearEnergyAtomicModel":
-        data = copy.deepcopy(data)
+        data = deepcopy(data)
         check_version_compatibility(data.pop("@version", 1), 2, 2)
         models = [
             BaseAtomicModel.get_class_by_type(model["type"]).deserialize(model)

deepmd/dpmodel/atomic_model/pairtab_atomic_model.py

@@ -1,5 +1,4 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-import copy
 from typing import (
     Optional,
     Union,
@@ -18,6 +17,9 @@
 from deepmd.dpmodel.utils.safe_gradient import (
     safe_for_sqrt,
 )
+from deepmd.utils.copy import (
+    deepcopy,
+)
 from deepmd.utils.pair_tab import (
     PairTab,
 )
@@ -174,7 +176,7 @@ def serialize(self) -> dict:
 
     @classmethod
     def deserialize(cls, data) -> "PairTabAtomicModel":
-        data = copy.deepcopy(data)
+        data = deepcopy(data)
         check_version_compatibility(data.pop("@version", 1), 2, 2)
         data.pop("@class")
         data.pop("type")

deepmd/dpmodel/descriptor/se_e2_a.py

@@ -1,5 +1,4 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-import copy
 import itertools
 from typing import (
     Any,
@@ -33,6 +32,9 @@
 from deepmd.env import (
     GLOBAL_NP_FLOAT_PRECISION,
 )
+from deepmd.utils.copy import (
+    deepcopy,
+)
 from deepmd.utils.data_system import (
     DeepmdDataSystem,
 )
@@ -460,7 +462,7 @@ def serialize(self) -> dict:
     @classmethod
     def deserialize(cls, data: dict) -> "DescrptSeA":
         """Deserialize from dict."""
-        data = copy.deepcopy(data)
+        data = deepcopy(data)
         check_version_compatibility(data.pop("@version", 1), 2, 1)
         data.pop("@class", None)
         data.pop("type", None)

deepmd/dpmodel/descriptor/se_r.py

@@ -1,5 +1,4 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-import copy
 from typing import (
     Any,
     Optional,
@@ -30,6 +29,9 @@
 from deepmd.dpmodel.utils.update_sel import (
     UpdateSel,
 )
+from deepmd.utils.copy import (
+    deepcopy,
+)
 from deepmd.utils.data_system import (
     DeepmdDataSystem,
 )
@@ -388,7 +390,7 @@ def serialize(self) -> dict:
     @classmethod
     def deserialize(cls, data: dict) -> "DescrptSeR":
         """Deserialize from dict."""
-        data = copy.deepcopy(data)
+        data = deepcopy(data)
         check_version_compatibility(data.pop("@version", 1), 2, 1)
         data.pop("@class", None)
         data.pop("type", None)

deepmd/dpmodel/descriptor/se_t.py

@@ -1,5 +1,4 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-import copy
 import itertools
 from typing import (
     Optional,
@@ -30,6 +29,9 @@
 from deepmd.dpmodel.utils.update_sel import (
     UpdateSel,
 )
+from deepmd.utils.copy import (
+    deepcopy,
+)
 from deepmd.utils.data_system import (
     DeepmdDataSystem,
 )
@@ -388,7 +390,7 @@ def serialize(self) -> dict:
     @classmethod
     def deserialize(cls, data: dict) -> "DescrptSeT":
         """Deserialize from dict."""
-        data = copy.deepcopy(data)
+        data = deepcopy(data)
         check_version_compatibility(data.pop("@version", 1), 2, 1)
         data.pop("@class", None)
         data.pop("type", None)

deepmd/dpmodel/fitting/dipole_fitting.py

@@ -1,5 +1,4 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-import copy
 from typing import (
     Any,
     Optional,
@@ -20,6 +19,9 @@
     OutputVariableDef,
     fitting_check_output,
 )
+from deepmd.utils.copy import (
+    deepcopy,
+)
 from deepmd.utils.version import (
     check_version_compatibility,
 )
@@ -156,7 +158,7 @@ def serialize(self) -> dict:
 
     @classmethod
     def deserialize(cls, data: dict) -> "GeneralFitting":
-        data = copy.deepcopy(data)
+        data = deepcopy(data)
         check_version_compatibility(data.pop("@version", 1), 2, 1)
         var_name = data.pop("var_name", None)
         assert var_name == "dipole"

deepmd/dpmodel/fitting/dos_fitting.py

@@ -1,5 +1,4 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-import copy
 from typing import (
     TYPE_CHECKING,
     Optional,
@@ -15,6 +14,9 @@
 from deepmd.dpmodel.fitting.invar_fitting import (
     InvarFitting,
 )
+from deepmd.utils.copy import (
+    deepcopy,
+)
 
 if TYPE_CHECKING:
     from deepmd.dpmodel.fitting.general_fitting import (
@@ -73,7 +75,7 @@ def __init__(
 
     @classmethod
     def deserialize(cls, data: dict) -> "GeneralFitting":
-        data = copy.deepcopy(data)
+        data = deepcopy(data)
         check_version_compatibility(data.pop("@version", 1), 2, 1)
         data["numb_dos"] = data.pop("dim_out")
         data.pop("tot_ener_zero", None)

deepmd/dpmodel/fitting/ener_fitting.py

@@ -1,5 +1,4 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-import copy
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -13,11 +12,15 @@
 from deepmd.dpmodel.fitting.invar_fitting import (
     InvarFitting,
 )
+from deepmd.utils.copy import (
+    deepcopy,
+)
 
 if TYPE_CHECKING:
     from deepmd.dpmodel.fitting.general_fitting import (
         GeneralFitting,
     )
+
 from deepmd.utils.version import (
     check_version_compatibility,
 )
@@ -73,7 +76,7 @@ def __init__(
 
     @classmethod
     def deserialize(cls, data: dict) -> "GeneralFitting":
-        data = copy.deepcopy(data)
+        data = deepcopy(data)
         check_version_compatibility(data.pop("@version", 1), 2, 1)
         data.pop("var_name")
         data.pop("dim_out")

deepmd/dpmodel/fitting/general_fitting.py

@@ -1,5 +1,4 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-import copy
 from abc import (
     abstractmethod,
 )
@@ -32,6 +31,9 @@
 from deepmd.env import (
     GLOBAL_NP_FLOAT_PRECISION,
 )
+from deepmd.utils.copy import (
+    deepcopy,
+)
 from deepmd.utils.finetune import (
     get_index_between_two_maps,
     map_atom_exclude_types,
@@ -320,7 +322,7 @@ def serialize(self) -> dict:
 
     @classmethod
     def deserialize(cls, data: dict) -> "GeneralFitting":
-        data = copy.deepcopy(data)
+        data = deepcopy(data)
         data.pop("@class")
         data.pop("type")
         variables = data.pop("@variables")

deepmd/dpmodel/fitting/invar_fitting.py

@@ -1,5 +1,4 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-import copy
 from typing import (
     Any,
     Optional,
@@ -16,6 +15,9 @@
     OutputVariableDef,
     fitting_check_output,
 )
+from deepmd.utils.copy import (
+    deepcopy,
+)
 from deepmd.utils.version import (
     check_version_compatibility,
 )
@@ -179,7 +181,7 @@ def serialize(self) -> dict:
 
     @classmethod
     def deserialize(cls, data: dict) -> "GeneralFitting":
-        data = copy.deepcopy(data)
+        data = deepcopy(data)
         check_version_compatibility(data.pop("@version", 1), 2, 1)
         return super().deserialize(data)
 

deepmd/dpmodel/fitting/polarizability_fitting.py

@@ -1,5 +1,4 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-import copy
 from typing import (
     Any,
     Optional,
@@ -26,6 +25,9 @@
     OutputVariableDef,
     fitting_check_output,
 )
+from deepmd.utils.copy import (
+    deepcopy,
+)
 from deepmd.utils.finetune import (
     get_index_between_two_maps,
 )
@@ -197,7 +199,7 @@ def serialize(self) -> dict:
 
     @classmethod
     def deserialize(cls, data: dict) -> "GeneralFitting":
-        data = copy.deepcopy(data)
+        data = deepcopy(data)
         check_version_compatibility(data.pop("@version", 1), 3, 1)
         var_name = data.pop("var_name", None)
         assert var_name == "polar"

deepmd/dpmodel/fitting/property_fitting.py

@@ -1,5 +1,4 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-import copy
 from typing import (
     Optional,
     Union,
@@ -13,6 +12,9 @@
 from deepmd.dpmodel.fitting.invar_fitting import (
     InvarFitting,
 )
+from deepmd.utils.copy import (
+    deepcopy,
+)
 from deepmd.utils.version import (
     check_version_compatibility,
 )
@@ -111,7 +113,7 @@ def __init__(
 
     @classmethod
     def deserialize(cls, data: dict) -> "PropertyFittingNet":
-        data = copy.deepcopy(data)
+        data = deepcopy(data)
         check_version_compatibility(data.pop("@version"), 2, 1)
         data.pop("dim_out")
         data.pop("var_name")

deepmd/dpmodel/utils/network.py

@@ -4,7 +4,6 @@
 See issue #2982 for more information.
 """
 
-import copy
 import itertools
 from typing import (
     Callable,
@@ -30,6 +29,9 @@
 from deepmd.dpmodel.utils.seed import (
     child_seed,
 )
+from deepmd.utils.copy import (
+    deepcopy,
+)
 from deepmd.utils.version import (
     check_version_compatibility,
 )
@@ -135,7 +137,7 @@ def deserialize(cls, data: dict) -> "NativeLayer":
         data : dict
             The dict to deserialize from.
         """
-        data = copy.deepcopy(data)
+        data = deepcopy(data)
         check_version_compatibility(data.pop("@version", 1), 1, 1)
         data.pop("@class", None)
         variables = data.pop("@variables")
@@ -404,7 +406,7 @@ def deserialize(cls, data: dict) -> "LayerNorm":
         data : dict
             The dict to deserialize from.
         """
-        data = copy.deepcopy(data)
+        data = deepcopy(data)
         check_version_compatibility(data.pop("@version", 1), 1, 1)
         data.pop("@class", None)
         variables = data.pop("@variables")
@@ -673,7 +675,7 @@ def deserialize(cls, data: dict) -> "EmbeddingNet":
             data : dict
                 The dict to deserialize from.
             """
-            data = copy.deepcopy(data)
+            data = deepcopy(data)
             check_version_compatibility(data.pop("@version", 1), 2, 1)
             data.pop("@class", None)
             layers = data.pop("layers")
@@ -778,7 +780,7 @@ def deserialize(cls, data: dict) -> "FittingNet":
             data : dict
                 The dict to deserialize from.
             """
-            data = copy.deepcopy(data)
+            data = deepcopy(data)
             check_version_compatibility(data.pop("@version", 1), 1, 1)
             data.pop("@class", None)
             layers = data.pop("layers")