fix bugs and enable more pd unitests

Author: HydrogenSulfate

deepmd/pd/infer/deep_eval.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.")
 
 

deepmd/pd/utils/env.py

@@ -135,6 +135,7 @@ def eval_model(
         if has_spin:
             input_dict["spin"] = batch_spin
         batch_output = model(**input_dict)
+        # 'atom_energy', 'energy', 'force', 'virial', 'mask'
         if isinstance(batch_output, tuple):
             batch_output = batch_output[0]
         if not return_tensor:

source/tests/pd/common.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,

source/tests/pd/model/test_forward_lower.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)

source/tests/pd/model/test_make_hessian_model.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,

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_denoise import (
     model_dpa1,
     model_dpa2,

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 import (  # model_dpau,
     model_dos,
     model_dpa1,

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_denoise import (
     model_dpa2,
 )

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_dpa1,
     model_dpa2,