fix bugs and enable more pd unitests

deepmodeling · Nov 1, 2024 · d03702a · d03702a
1 parent 26047e9
commit d03702a
Show file tree

Hide file tree

Showing 10 changed files with 158 additions and 254 deletions.
diff --git a/deepmd/pd/infer/deep_eval.py b/deepmd/pd/infer/deep_eval.py
diff --git a/deepmd/pd/utils/env.py b/deepmd/pd/utils/env.py
@@ -15,8 +15,6 @@
     set_default_nthreads,
 )
 
-log = logging.getLogger(__name__)
-
 SAMPLER_RECORD = os.environ.get("SAMPLER_RECORD", False)
 try:
     # only linux
@@ -87,6 +85,7 @@ def enable_prim(enable: bool = True):
 
         core.set_prim_eager_enabled(True)
         core._set_prim_all_enabled(True)
+        log = logging.getLogger(__name__)
         log.info("Enable prim in eager and static mode.")
 
 

diff --git a/source/tests/pd/model/test_forward_lower.py b/source/tests/pd/model/test_forward_lower.py
@@ -5,9 +5,6 @@
 import numpy as np
 import paddle
 
-from deepmd.pd.infer.deep_eval import (
-    eval_model,
-)
 from deepmd.pd.model.model import (
     get_model,
 )
@@ -22,6 +19,9 @@
 from ...seed import (
     GLOBAL_SEED,
 )
+from ..common import (
+    eval_model,
+)
 from .test_permutation import (  # model_dpau,
     model_dpa1,
     model_dpa2,

diff --git a/source/tests/pd/model/test_make_hessian_model.py b/source/tests/pd/model/test_make_hessian_model.py
@@ -137,7 +137,7 @@ def ff(xx):
         )
 
 
-@unittest.skip("TODO")
+@unittest.skip("Skip temporarily")
 class TestDPModel(unittest.TestCase, HessianTest):
     def setUp(self):
         paddle.seed(2)

diff --git a/source/tests/pd/model/test_permutation_denoise.py b/source/tests/pd/model/test_permutation_denoise.py
@@ -5,9 +5,6 @@
 import numpy as np
 import paddle
 
-from deepmd.pd.infer.deep_eval import (
-    eval_model,
-)
 from deepmd.pd.model.model import (
     get_model,
 )
@@ -18,6 +15,9 @@
 from ...seed import (
     GLOBAL_SEED,
 )
+from ..common import (
+    eval_model,
+)
 from .test_permutation import (  # model_dpau,
     model_dpa1,
     model_dpa2,

diff --git a/source/tests/pd/model/test_rot_denoise.py b/source/tests/pd/model/test_rot_denoise.py
@@ -5,9 +5,6 @@
 import numpy as np
 import paddle
 
-from deepmd.pd.infer.deep_eval import (
-    eval_model,
-)
 from deepmd.pd.model.model import (
     get_model,
 )
@@ -18,6 +15,9 @@
 from ...seed import (
     GLOBAL_SEED,
 )
+from ..common import (
+    eval_model,
+)
 from .test_permutation_denoise import (
     model_dpa1,
     model_dpa2,

diff --git a/source/tests/pd/model/test_smooth.py b/source/tests/pd/model/test_smooth.py
@@ -5,9 +5,6 @@
 import numpy as np
 import paddle
 
-from deepmd.pd.infer.deep_eval import (
-    eval_model,
-)
 from deepmd.pd.model.model import (
     get_model,
 )
@@ -18,6 +15,9 @@
 from ...seed import (
     GLOBAL_SEED,
 )
+from ..common import (
+    eval_model,
+)
 from .test_permutation import (  # model_dpau,
     model_dos,
     model_dpa1,

diff --git a/source/tests/pd/model/test_smooth_denoise.py b/source/tests/pd/model/test_smooth_denoise.py
@@ -5,9 +5,6 @@
 import numpy as np
 import paddle
 
-from deepmd.pd.infer.deep_eval import (
-    eval_model,
-)
 from deepmd.pd.model.model import (
     get_model,
 )
@@ -18,6 +15,9 @@
 from ...seed import (
     GLOBAL_SEED,
 )
+from ..common import (
+    eval_model,
+)
 from .test_permutation_denoise import (
     model_dpa2,
 )

diff --git a/source/tests/pd/model/test_trans_denoise.py b/source/tests/pd/model/test_trans_denoise.py
@@ -5,9 +5,6 @@
 import numpy as np
 import paddle
 
-from deepmd.pd.infer.deep_eval import (
-    eval_model,
-)
 from deepmd.pd.model.model import (
     get_model,
 )
@@ -18,6 +15,9 @@
 from ...seed import (
     GLOBAL_SEED,
 )
+from ..common import (
+    eval_model,
+)
 from .test_permutation_denoise import (
     model_dpa1,
     model_dpa2,

diff --git a/source/tests/pd/model/test_unused_params.py b/source/tests/pd/model/test_unused_params.py
@@ -4,9 +4,6 @@
 
 import paddle
 
-from deepmd.pd.infer.deep_eval import (
-    eval_model,
-)
 from deepmd.pd.model.model import (
     get_model,
 )
@@ -17,6 +14,9 @@
 from ...seed import (
     GLOBAL_SEED,
 )
+from ..common import (
+    eval_model,
+)
 from .test_permutation import (
     model_dpa2,
 )