Cache TF Versions (awslabs#421)

Author: NihalHarish

smdebug/pytorch/utils.py

@@ -7,6 +7,9 @@
 from smdebug.core.reduction_config import ALLOWED_NORMS, ALLOWED_REDUCTIONS
 from smdebug.core.reductions import get_numpy_reduction
 
+# Cached Pytorch Version
+PT_VERSION = version.parse(torch.__version__)
+
 
 def get_reduction_of_data(reduction_name, tensor_data, tensor_name, abs=False):
     if isinstance(tensor_data, np.ndarray):
@@ -42,20 +45,20 @@ def is_pt_1_5():
     Determine whether the version of torch is 1.5.x
     :return: bool
     """
-    return version.parse("1.5.0") <= version.parse(torch.__version__) < version.parse("1.6.0")
+    return version.parse("1.5.0") <= PT_VERSION < version.parse("1.6.0")
 
 
 def is_pt_1_6():
     """
     Determine whether the version of torch is 1.6.x
     :return: bool
     """
-    return version.parse("1.6.0") <= version.parse(torch.__version__) < version.parse("1.7.0")
+    return version.parse("1.6.0") <= PT_VERSION < version.parse("1.7.0")
 
 
 def is_pt_1_7():
     """
     Determine whether the version of torch is 1.7.x
     :return: bool
     """
-    return version.parse("1.7.0") <= version.parse(torch.__version__) < version.parse("1.8.0")
+    return version.parse("1.7.0") <= PT_VERSION < version.parse("1.8.0")

smdebug/tensorflow/session.py

@@ -13,6 +13,7 @@
 from .base_hook import TensorflowBaseHook
 from .tensor_ref import TensorType
 from .utils import (
+    TF_VERSION,
     TFDistributionStrategy,
     build_fetches_tuple,
     extract_graph_summary,
@@ -241,7 +242,7 @@ def _is_not_supported(self):
             if self.distribution_strategy == TFDistributionStrategy.MIRRORED:
                 from packaging import version
 
-                if version.parse(tf.__version__) < version.parse("1.14.0"):
+                if TF_VERSION < version.parse("1.14.0"):
                     self._hook_supported = False
                     # in tf 1.13, we can't support mirrored strategy as
                     # MirroredVariable does not have _values attribute

smdebug/tensorflow/utils.py

@@ -13,11 +13,14 @@
 # First Party
 from smdebug.core.modes import ModeKeys
 
+# Cached TF Version
+TF_VERSION = version.parse(tf.__version__)
+
 
 def does_tf_support_mixed_precision_training():
     # The Keras mixed precision API is first available in TensorFlow 2.1.0
     # See: https://www.tensorflow.org/guide/mixed_precision
-    return version.parse(tf.__version__) >= version.parse("2.1.0")
+    return TF_VERSION >= version.parse("2.1.0")
 
 
 def supported_tf_variables():
@@ -405,23 +408,23 @@ def get_keras_mode(mode):
 
 
 def is_tf_version_2x():
-    return version.parse(tf.__version__) >= version.parse("2.0.0")
+    return TF_VERSION >= version.parse("2.0.0")
 
 
 def is_tf_version_2_2_x():
-    return version.parse("2.2.0") <= version.parse(tf.__version__) < version.parse("2.3.0")
+    return version.parse("2.2.0") <= TF_VERSION < version.parse("2.3.0")
 
 
 def is_tf_version_2_3_x():
-    return version.parse("2.3.0") <= version.parse(tf.__version__) < version.parse("2.4.0")
+    return version.parse("2.3.0") <= TF_VERSION < version.parse("2.4.0")
 
 
 def is_tf_version_2_4_x():
-    return version.parse("2.4.0") <= version.parse(tf.__version__) < version.parse("2.5.0")
+    return version.parse("2.4.0") <= TF_VERSION < version.parse("2.5.0")
 
 
 def is_tf_version_greater_than_2_4_x():
-    return version.parse("2.4.0") <= version.parse(tf.__version__)
+    return version.parse("2.4.0") <= TF_VERSION
 
 
 def is_profiler_supported_for_tf_version():

tests/tensorflow2/utils.py

@@ -3,6 +3,9 @@
 import tensorflow.compat.v2 as tf
 from packaging import version
 
+# Cached TF Version
+TF_VERSION = version.parse(tf.__version__)
+
 
 def is_tf_2_2():
     """
@@ -12,12 +15,12 @@ def is_tf_2_2():
     number of tensor_names emitted by 1.
     :return: bool
     """
-    if version.parse(tf.__version__) >= version.parse("2.2.0"):
+    if TF_VERSION >= version.parse("2.2.0"):
         return True
     return False
 
 
 def is_tf_2_3():
-    if version.parse(tf.__version__) == version.parse("2.3.0"):
+    if TF_VERSION == version.parse("2.3.0"):
         return True
     return False