tests: move init_models to setUpModule (deepmodeling#3715)

## Summary by CodeRabbit - **Refactor** - Enhanced the initialization process of global variables across multiple test modules to improve test setup efficiency and consistency.  --------- Signed-off-by: Jinzhe Zeng <[email protected]> Co-authored-by: Han Wang <[email protected]>
mtaillefumier · Sep 18, 2024 · 2f65b08 · 2f65b08
1 parent f734015
commit 2f65b08
Show file tree

Hide file tree

Showing 9 changed files with 58 additions and 12 deletions.
diff --git a/deepmd/tf/common.py b/deepmd/tf/common.py
@@ -13,6 +13,9 @@
 )
 
 import tensorflow
+from packaging.version import (
+    Version,
+)
 from tensorflow.python.framework import (
     tensor_util,
 )
@@ -31,6 +34,7 @@
 )
 from deepmd.tf.env import (
     GLOBAL_TF_FLOAT_PRECISION,
+    TF_VERSION,
     op_module,
     tf,
 )
@@ -289,3 +293,8 @@ def clear_session():
     tf.reset_default_graph()
     # TODO: remove this line when data_requirement is not a global variable
     data_requirement.clear()
+    _TF_VERSION = Version(TF_VERSION)
+    if _TF_VERSION < Version("2.4.0"):
+        tf.train.experimental.disable_mixed_precision_graph_rewrite()
+    else:
+        tf.mixed_precision.disable_mixed_precision_graph_rewrite()
diff --git a/source/tests/pt/test_multitask.py b/source/tests/pt/test_multitask.py
@@ -25,9 +25,12 @@
     model_se_e2_a,
 )
 
-multitask_template_json = str(Path(__file__).parent / "water/multitask.json")
-with open(multitask_template_json) as f:
-    multitask_template = json.load(f)
+
+def setUpModule():
+    global multitask_template
+    multitask_template_json = str(Path(__file__).parent / "water/multitask.json")
+    with open(multitask_template_json) as f:
+        multitask_template = json.load(f)
 
 
 class MultiTaskTrainTest:

diff --git a/source/tests/tf/test_mixed_prec_training.py b/source/tests/tf/test_mixed_prec_training.py
@@ -8,6 +8,9 @@
     Version,
 )
 
+from deepmd.tf.common import (
+    clear_session,
+)
 from deepmd.tf.env import (
     TF_VERSION,
 )
@@ -61,3 +64,4 @@ def tearDown(self):
         _file_delete("model.ckpt-100.data-00000-of-00001")
         _file_delete("input_v2_compat.json")
         _file_delete("lcurve.out")
+        clear_session()
diff --git a/source/tests/tf/test_model_compression_se_a.py b/source/tests/tf/test_model_compression_se_a.py
@@ -73,8 +73,16 @@ def _init_models_exclude_types():
     return INPUT, frozen_model, compressed_model
 
 
-INPUT, FROZEN_MODEL, COMPRESSED_MODEL = _init_models()
-INPUT_ET, FROZEN_MODEL_ET, COMPRESSED_MODEL_ET = _init_models_exclude_types()
+def setUpModule():
+    global \
+        INPUT, \
+        FROZEN_MODEL, \
+        COMPRESSED_MODEL, \
+        INPUT_ET, \
+        FROZEN_MODEL_ET, \
+        COMPRESSED_MODEL_ET
+    INPUT, FROZEN_MODEL, COMPRESSED_MODEL = _init_models()
+    INPUT_ET, FROZEN_MODEL_ET, COMPRESSED_MODEL_ET = _init_models_exclude_types()
 
 
 class TestDeepPotAPBC(unittest.TestCase):

diff --git a/source/tests/tf/test_model_compression_se_a_ebd.py b/source/tests/tf/test_model_compression_se_a_ebd.py
@@ -85,8 +85,16 @@ def _init_models_exclude_types():
     return INPUT, frozen_model, compressed_model
 
 
-INPUT, FROZEN_MODEL, COMPRESSED_MODEL = _init_models()
-INPUT_ET, FROZEN_MODEL_ET, COMPRESSED_MODEL_ET = _init_models_exclude_types()
+def setUpModule():
+    global \
+        INPUT, \
+        FROZEN_MODEL, \
+        COMPRESSED_MODEL, \
+        INPUT_ET, \
+        FROZEN_MODEL_ET, \
+        COMPRESSED_MODEL_ET
+    INPUT, FROZEN_MODEL, COMPRESSED_MODEL = _init_models()
+    INPUT_ET, FROZEN_MODEL_ET, COMPRESSED_MODEL_ET = _init_models_exclude_types()
 
 
 class TestDeepPotAPBC(unittest.TestCase):

diff --git a/source/tests/tf/test_model_compression_se_a_ebd_type_one_side.py b/source/tests/tf/test_model_compression_se_a_ebd_type_one_side.py
@@ -85,7 +85,9 @@ def _init_models_exclude_types():
     return INPUT, frozen_model, compressed_model
 
 
-INPUT, FROZEN_MODEL, COMPRESSED_MODEL = _init_models()
+def setUpModule():
+    global INPUT, FROZEN_MODEL, COMPRESSED_MODEL
+    INPUT, FROZEN_MODEL, COMPRESSED_MODEL = _init_models()
 
 
 class TestDeepPotAPBC(unittest.TestCase):

diff --git a/source/tests/tf/test_model_compression_se_atten.py b/source/tests/tf/test_model_compression_se_atten.py
@@ -154,8 +154,16 @@ def _init_models_exclude_types():
     return inputs, frozen_models, compressed_models
 
 
-INPUTS, FROZEN_MODELS, COMPRESSED_MODELS = _init_models()
-INPUTS_ET, FROZEN_MODELS_ET, COMPRESSED_MODELS_ET = _init_models_exclude_types()
+def setUpModule():
+    global \
+        INPUTS, \
+        FROZEN_MODELS, \
+        COMPRESSED_MODELS, \
+        INPUTS_ET, \
+        FROZEN_MODELS_ET, \
+        COMPRESSED_MODELS_ET
+    INPUTS, FROZEN_MODELS, COMPRESSED_MODELS = _init_models()
+    INPUTS_ET, FROZEN_MODELS_ET, COMPRESSED_MODELS_ET = _init_models_exclude_types()
 
 
 @unittest.skipIf(

diff --git a/source/tests/tf/test_model_compression_se_r.py b/source/tests/tf/test_model_compression_se_r.py
@@ -60,7 +60,9 @@ def _init_models():
     return INPUT, frozen_model, compressed_model
 
 
-INPUT, FROZEN_MODEL, COMPRESSED_MODEL = _init_models()
+def setUpModule():
+    global INPUT, FROZEN_MODEL, COMPRESSED_MODEL
+    INPUT, FROZEN_MODEL, COMPRESSED_MODEL = _init_models()
 
 
 class TestDeepPotAPBC(unittest.TestCase):

diff --git a/source/tests/tf/test_model_compression_se_t.py b/source/tests/tf/test_model_compression_se_t.py
@@ -60,7 +60,9 @@ def _init_models():
     return INPUT, frozen_model, compressed_model
 
 
-INPUT, FROZEN_MODEL, COMPRESSED_MODEL = _init_models()
+def setUpModule():
+    global INPUT, FROZEN_MODEL, COMPRESSED_MODEL
+    INPUT, FROZEN_MODEL, COMPRESSED_MODEL = _init_models()
 
 
 def tearDownModule():