disable_mixed_precision_graph_rewrite

Signed-off-by: Jinzhe Zeng <[email protected]>
deepmodeling · May 6, 2024 · e4e07ce · e4e07ce
1 parent ecfa478
commit e4e07ce
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 0 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/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()