Use packaging library to avoid DeprecationWarning from distutils

Signed-off-by: Deyu Huang <[email protected]>

Author: hwangdeyu

tests/common.py

@@ -9,7 +9,7 @@
 import unittest
 from collections import defaultdict
 
-from distutils.version import LooseVersion
+from packaging.version import Version
 from parameterized import parameterized
 import numpy as np
 import tensorflow as tf
@@ -98,7 +98,7 @@ def _get_backend_version(self):
             pass
 
         if version:
-            version = LooseVersion(version)
+            version = Version(version)
         return version
 
     def __str__(self):
@@ -178,7 +178,7 @@ def check_opset_after_tf_version(tf_version, required_opset, message=""):
     """ Skip if tf_version > max_required_version """
     config = get_test_config()
     reason = _append_message("conversion requires opset {} after tf {}".format(required_opset, tf_version), message)
-    skip = config.tf_version >= LooseVersion(tf_version) and config.opset < required_opset
+    skip = config.tf_version >= Version(tf_version) and config.opset < required_opset
     return unittest.skipIf(skip, reason)
 
 
@@ -284,7 +284,7 @@ def check_tfjs_max_version(max_accepted_version, message=""):
     except ModuleNotFoundError:
         can_import = False
     return unittest.skipIf(can_import and not config.skip_tfjs_tests and \
-         tensorflowjs.__version__ > LooseVersion(max_accepted_version), reason)
+         Version(tensorflowjs.__version__) > Version(max_accepted_version), reason)
 
 def check_tfjs_min_version(min_required_version, message=""):
     """ Skip if tjs_version < min_required_version """
@@ -296,20 +296,20 @@ def check_tfjs_min_version(min_required_version, message=""):
     except ModuleNotFoundError:
         can_import = False
     return unittest.skipIf(can_import and not config.skip_tfjs_tests and \
-         tensorflowjs.__version__ < LooseVersion(min_required_version), reason)
+         Version(tensorflowjs.__version__) < Version(min_required_version), reason)
 
 def check_tf_max_version(max_accepted_version, message=""):
     """ Skip if tf_version > max_required_version """
     config = get_test_config()
     reason = _append_message("conversion requires tf <= {}".format(max_accepted_version), message)
-    return unittest.skipIf(config.tf_version > LooseVersion(max_accepted_version), reason)
+    return unittest.skipIf(config.tf_version > Version(max_accepted_version), reason)
 
 
 def check_tf_min_version(min_required_version, message=""):
     """ Skip if tf_version < min_required_version """
     config = get_test_config()
     reason = _append_message("conversion requires tf >= {}".format(min_required_version), message)
-    return unittest.skipIf(config.tf_version < LooseVersion(min_required_version), reason)
+    return unittest.skipIf(config.tf_version < Version(min_required_version), reason)
 
 
 def skip_tf_versions(excluded_versions, message=""):
@@ -385,7 +385,7 @@ def check_onnxruntime_min_version(min_required_version, message=""):
     config = get_test_config()
     reason = _append_message("conversion requires onnxruntime >= {}".format(min_required_version), message)
     return unittest.skipIf(config.is_onnxruntime_backend and
-                           config.backend_version < LooseVersion(min_required_version), reason)
+                           config.backend_version < Version(min_required_version), reason)
 
 
 def skip_caffe2_backend(message=""):

tests/run_pretrained_models.py

@@ -17,7 +17,7 @@
 import zipfile
 import random
 from collections import namedtuple
-from distutils.version import LooseVersion
+from packaging.version import Version
 
 
 import yaml
@@ -789,7 +789,7 @@ def main():
                 continue
 
             if t.tf_min_version:
-                if tf_utils.get_tf_version() < LooseVersion(str(t.tf_min_version)):
+                if tf_utils.get_tf_version() < Version(str(t.tf_min_version)):
                     logger.info("Skip %s: %s %s", test, "Min TF version needed:", t.tf_min_version)
                     continue
 

tests/test_backend.py

@@ -5,11 +5,11 @@
 
 import os
 import unittest
-from packaging.version import Version
 from itertools import product
 
 import numpy as np
 from numpy.testing import assert_almost_equal
+from packaging.version import Version
 import tensorflow as tf
 
 from tensorflow.python.ops import lookup_ops

tf2onnx/convert.py

@@ -10,7 +10,7 @@
 import argparse
 import os
 import sys
-from distutils.version import LooseVersion
+from packaging.version import Version
 
 os.environ['TF_CPP_MIN_LOG_LEVEL'] = "3"
 
@@ -20,7 +20,7 @@
 from tf2onnx import constants, logging, utils, optimizer
 from tf2onnx import tf_loader
 from tf2onnx.graph import ExternalTensorStorage
-from tf2onnx.tf_utils import compress_graph_def
+from tf2onnx.tf_utils import compress_graph_def, get_tf_version
 
 
 
@@ -431,7 +431,7 @@ def from_keras(model, input_signature=None, opset=None, custom_ops=None, custom_
     Returns:
         An ONNX model_proto and an external_tensor_storage dict.
     """
-    if LooseVersion(tf.__version__) < "2.0":
+    if get_tf_version() < Version("2.0"):
         return _from_keras_tf1(model, opset, custom_ops, custom_op_handlers, custom_rewriter, inputs_as_nchw,
                                outputs_as_nchw, extra_opset, shape_override, target, large_model, output_path)
 
@@ -540,7 +540,7 @@ def from_function(function, input_signature=None, opset=None, custom_ops=None, c
     Returns:
         An ONNX model_proto and an external_tensor_storage dict.
     """
-    if LooseVersion(tf.__version__) < "2.0":
+    if get_tf_version() < Version("2.0"):
         raise NotImplementedError("from_function requires tf-2.0 or newer")
 
     if input_signature is None:

tf2onnx/shape_inference.py

@@ -6,9 +6,9 @@
 """
 
 import logging
-from distutils.version import LooseVersion
 from collections import defaultdict
 import numpy as np
+from packaging.version import Version
 from tf2onnx import utils
 from tf2onnx.tf_utils import get_tf_tensor_shape, get_tf_const_value, get_tf_shape_attr, get_tf_version
 from tf2onnx.tf_loader import tf_reload_graph
@@ -32,7 +32,7 @@ def infer_shape(tf_graph, shape_override):
 
     op_outputs_with_none_shape = check_shape_for_tf_graph(tf_graph)
     if op_outputs_with_none_shape:
-        if get_tf_version() > LooseVersion("1.5.0"):
+        if get_tf_version() > Version("1.5.0"):
             for op, outs in op_outputs_with_none_shape.items():
                 logger.warning(
                     "Cannot infer shape for %s: %s",

tf2onnx/tf_loader.py

@@ -5,7 +5,7 @@
 
 import logging
 import uuid
-from distutils.version import LooseVersion
+from packaging.version import Version
 
 import tensorflow as tf
 import numpy as np
@@ -75,7 +75,7 @@ def not_implemented_tf_placeholder(*args, **kwargs):
     tf_placeholder = tf.compat.v1.placeholder
     tf_placeholder_with_default = tf.compat.v1.placeholder_with_default
     extract_sub_graph = tf.compat.v1.graph_util.extract_sub_graph
-elif LooseVersion(tf.__version__) >= "1.13":
+elif Version(tf.__version__) >= Version("1.13"):
     # 1.13 introduced the compat namespace
     tf_reset_default_graph = tf.compat.v1.reset_default_graph
     tf_global_variables = tf.compat.v1.global_variables
@@ -162,7 +162,7 @@ def make_tensor_proto_wrapped(values, dtype=None, shape=None, verify_shape=False
 
     try:
         function_converter = _FunctionConverterData
-        if LooseVersion(tf.__version__) >= "2.6.0":
+        if Version(tf.__version__) >= Version("2.6.0"):
             from tensorflow.python.eager import context
             from tensorflow.python.framework.convert_to_constants import _FunctionConverterDataInEager, \
                 _FunctionConverterDataInGraph
@@ -267,7 +267,7 @@ def from_function(func, input_names, output_names, large_model=False):
         return convert_variables_to_constants_large_model(func)
 
     try:
-        if get_tf_version() < LooseVersion("2.2"):
+        if get_tf_version() < Version("2.2"):
             frozen_func = convert_variables_to_constants_v2(func, lower_control_flow=False)
         else:
             frozen_func = convert_variables_to_constants_v2(func, lower_control_flow=False, aggressive_inlining=True)
@@ -687,7 +687,7 @@ def tf_optimize_grappler(input_names, output_names, graph_def):
         'constfold', 'function'
     ]
 
-    if LooseVersion(tf.__version__) >= "2.5":
+    if Version(tf.__version__) >= Version("2.5"):
         # This flag disables folding QDQ nodes around constants in the network (eg: around conv/FC weights)
         rewrite_options.experimental_disable_folding_quantization_emulation = True
 
@@ -710,7 +710,7 @@ def tf_optimize(input_names, output_names, graph_def):
                    [utils.node_name(i) for i in output_names]
     graph_def = extract_sub_graph(graph_def, needed_names)
 
-    want_grappler = is_tf2() or LooseVersion(tf.__version__) >= "1.15"
+    want_grappler = is_tf2() or Version(tf.__version__) >= Version("1.15")
     if want_grappler:
         graph_def = tf_optimize_grappler(input_names, output_names, graph_def)
     else:
@@ -730,7 +730,7 @@ def tf_optimize(input_names, output_names, graph_def):
 def tf_reload_graph(tf_graph):
     """Invoke tensorflow cpp shape inference by reloading graph_def."""
     # invoke c api if tf version is below 1.8
-    if get_tf_version() < LooseVersion("1.8"):
+    if get_tf_version() < Version("1.8"):
         logger.debug(
             "On TF < 1.8, graph is constructed by python API, "
             "which doesn't invoke shape inference, please set "

tf2onnx/tf_utils.py

@@ -6,7 +6,7 @@
 """
 
 import collections
-from distutils.version import LooseVersion
+from packaging.version import Version
 
 import numpy as np
 import tensorflow as tf
@@ -121,7 +121,7 @@ def get_tf_node_attr(node, name):
 
 
 def get_tf_version():
-    return LooseVersion(tf.__version__)
+    return Version(tf.__version__)
 
 def compress_graph_def(graph_def):
     """