Support TF GNR-BASE (#1415)

Signed-off-by: zehao-intel <zehao.huang@intel.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
intel · Nov 28, 2023 · 8160c71 · 8160c71
1 parent d131832
commit 8160c71
Show file tree

Hide file tree

Showing 10 changed files with 37 additions and 9 deletions.
diff --git a/examples/tensorflow/oob_models/quantization/ptq/tf_benchmark.py b/examples/tensorflow/oob_models/quantization/ptq/tf_benchmark.py
@@ -94,6 +94,8 @@ def create_tf_config(args):
     config.allow_soft_placement = True
     config.intra_op_parallelism_threads = OMP_NUM_THREADS
     config.inter_op_parallelism_threads = 1
+    # additional options
+    config.graph_options.rewrite_options.function_optimization = rewriter_config_pb2.RewriterConfig.AGGRESSIVE
     if args.precision == 'bfloat16':
         config.graph_options.rewrite_options.auto_mixed_precision_mkl = rewriter_config_pb2.RewriterConfig.ON
     return config

diff --git a/neural_compressor/adaptor/tensorflow.py b/neural_compressor/adaptor/tensorflow.py
@@ -44,7 +44,15 @@
 from .query import QueryBackendCapability
 
 tensorflow = LazyImport("tensorflow")
-spr_base_verions = ("2.11.0202242", "2.11.0202250", "2.11.0202317", "2.11.0202323")
+spr_base_verions = (
+    "2.11.0202242",
+    "2.11.0202250",
+    "2.11.0202317",
+    "2.11.0202323",
+    "2.14.0202335",
+    "2.14.dev202335",
+    "2.15.0202341",
+)
 
 
 @adaptor_registry
@@ -2146,7 +2154,15 @@ def _compare(version1, version2):
             if self.version in sub_data["version"]["name"]:
                 return sub_data
             else:
-                if sub_data["version"]["name"] == ["2.11.0202242", "2.11.0202250", "2.11.0202317", "2.11.0202323"]:
+                if sub_data["version"]["name"] == [
+                    "2.11.0202242",
+                    "2.11.0202250",
+                    "2.11.0202317",
+                    "2.11.0202323",
+                    "2.14.0202335",
+                    "2.14.dev202335",
+                    "2.15.0202341",
+                ]:
                     continue
                 sorted_list = copy.deepcopy(sub_data["version"]["name"])
                 sorted_list.remove("default") if "default" in sorted_list else None

diff --git a/neural_compressor/adaptor/tensorflow.yaml b/neural_compressor/adaptor/tensorflow.yaml
@@ -16,7 +16,7 @@
 ---
 -
   version:
-    name: ['2.11.0202242', '2.11.0202250', '2.11.0202317', '2.11.0202323']
+    name: ['2.11.0202242', '2.11.0202250', '2.11.0202317', '2.11.0202323', '2.14.0202335', '2.14.dev202335', '2.15.0202341']
 
   bf16: ["_MklLayerNorm", "Conv2D", "Conv2DBackpropFilter", "Conv2DBackpropInput", "Conv3D", "Conv3DBackpropFilterV2", "Conv3DBackpropInputV2",
           "DepthwiseConv2dNative", "DepthwiseConv2dNativeBackpropFilter", "DepthwiseConv2dNativeBackpropInput", "GRUBlockCell",

diff --git a/neural_compressor/adaptor/tf_utils/graph_rewriter/bf16/bf16_convert.py b/neural_compressor/adaptor/tf_utils/graph_rewriter/bf16/bf16_convert.py
@@ -31,9 +31,9 @@
 from neural_compressor.adaptor.tf_utils.graph_util import GraphRewriterHelper as Helper
 from neural_compressor.adaptor.tf_utils.util import TF_SPR_BASE_VERSIONS
 
-from ..generic.dequantize_cast_optimizer import DequantizeCastOptimizer
 from ..generic.graph_cse_optimizer import GraphCseOptimizer
 from ..graph_base import GraphRewriterBase
+from .dequantize_cast_optimizer import DequantizeCastOptimizer
 
 DT_FLOAT32 = attr_value_pb2.AttrValue(type=dtypes.float32.as_datatype_enum)
 DT_BFLOAT16 = attr_value_pb2.AttrValue(type=dtypes.bfloat16.as_datatype_enum)

diff --git a/...iter/generic/dequantize_cast_optimizer.py → ...ewriter/bf16/dequantize_cast_optimizer.py b/...iter/generic/dequantize_cast_optimizer.py → ...ewriter/bf16/dequantize_cast_optimizer.py
diff --git a/neural_compressor/adaptor/tf_utils/graph_rewriter/generic/convert_add_to_biasadd.py b/neural_compressor/adaptor/tf_utils/graph_rewriter/generic/convert_add_to_biasadd.py
@@ -21,6 +21,7 @@
 
 from neural_compressor.adaptor.tf_utils.graph_util import GraphAnalyzer
 from neural_compressor.adaptor.tf_utils.graph_util import GraphRewriterHelper as Helper
+from neural_compressor.adaptor.tf_utils.util import TF_SPR_BASE_VERSIONS
 from neural_compressor.utils.utility import dump_elapsed_time
 
 from ..graph_base import GraphRewriterBase
@@ -38,7 +39,7 @@ def do_transformation(self):
 
         import tensorflow as tf
 
-        if tf.version.VERSION not in ("2.11.0202242", "2.11.0202250", "2.11.0202317", "2.11.0202323"):
+        if tf.version.VERSION not in TF_SPR_BASE_VERSIONS:
             target_nodes = g.query_fusion_pattern_nodes([["MatMul", "Conv2D"], ["Add", "AddV2"]])
         else:
             target_nodes = g.query_fusion_pattern_nodes([["MatMul"], ["Add", "AddV2"]])

diff --git a/neural_compressor/adaptor/tf_utils/util.py b/neural_compressor/adaptor/tf_utils/util.py
@@ -38,7 +38,15 @@
 
 from .graph_util import GraphAnalyzer, GraphRewriterHelper
 
-TF_SPR_BASE_VERSIONS = ("2.11.0202242", "2.11.0202250", "2.11.0202317", "2.11.0202323")
+TF_SPR_BASE_VERSIONS = (
+    "2.11.0202242",
+    "2.11.0202250",
+    "2.11.0202317",
+    "2.11.0202323",
+    "2.14.0202335",
+    "2.14.dev202335",
+    "2.15.0202341",
+)
 
 
 def version1_lt_version2(version1, version2):

diff --git a/neural_compressor/model/tensorflow_model.py b/neural_compressor/model/tensorflow_model.py
@@ -81,7 +81,8 @@ def get_model_type(model):
         return "graph"
     elif isinstance(model, tf.compat.v1.GraphDef):
         return "graph_def"
-    elif isinstance(model, tf.compat.v1.estimator.Estimator):
+    # tf.estimator is removed after tf2.14.0
+    elif version1_lt_version2(tf.__version__, "2.14.0") and isinstance(model, tf.compat.v1.estimator.Estimator):
         return "estimator"
     elif isinstance(model, str):
         model = os.path.abspath(os.path.expanduser(model))

diff --git a/test/adaptor/tensorflow_adaptor/test_tensorflow_graph_dequantize_cast_optimizer.py b/test/adaptor/tensorflow_adaptor/test_tensorflow_graph_dequantize_cast_optimizer.py
@@ -6,7 +6,7 @@
 import yaml
 from tensorflow.python.framework import dtypes
 
-from neural_compressor.adaptor.tf_utils.graph_rewriter.generic.dequantize_cast_optimizer import DequantizeCastOptimizer
+from neural_compressor.adaptor.tf_utils.graph_rewriter.bf16.dequantize_cast_optimizer import DequantizeCastOptimizer
 from neural_compressor.adaptor.tf_utils.graph_util import GraphRewriterHelper as Helper
 from neural_compressor.adaptor.tf_utils.util import disable_random
 

diff --git a/test/tfnewapi/test_tensorflow_graph_dequantize_cast_optimizer_newapi.py b/test/tfnewapi/test_tensorflow_graph_dequantize_cast_optimizer_newapi.py
@@ -6,7 +6,7 @@
 import yaml
 from tensorflow.python.framework import dtypes
 
-from neural_compressor.adaptor.tf_utils.graph_rewriter.generic.dequantize_cast_optimizer import DequantizeCastOptimizer
+from neural_compressor.adaptor.tf_utils.graph_rewriter.bf16.dequantize_cast_optimizer import DequantizeCastOptimizer
 from neural_compressor.adaptor.tf_utils.graph_util import GraphRewriterHelper as Helper
 from neural_compressor.adaptor.tf_utils.util import disable_random