Fix graph_optimization bf16 convert issue (#1306)

Author: lvliang-intel

neural_compressor/experimental/graph_optimization.py

@@ -126,9 +126,14 @@ def __call__(self):
             self._model.output_tensor_names = cfg.model.outputs
             self._model.input_tensor_names = cfg.model.inputs
             self._model.workspace_path = cfg.tuning.workspace.path
-        else: 
-            logger.warning("Only TensorFlow graph optimization is supported at current stage.") 
-            sys.exit(0) 
+
+            if 'bf16' in self._precisions or \
+               (cfg.mixed_precision and 'bf16' in cfg.mixed_precision.precisions) or \
+               (cfg.graph_optimization and 'bf16' in cfg.graph_optimization.precisions):
+                os.environ['MIX_PRECISION_TEST'] = '1'
+        else:
+            logger.warning("Only TensorFlow graph optimization is supported at current stage.")
+            sys.exit(0)
 
         # when eval_func is set, will be directly used and eval_dataloader can be None
         if self._eval_func is None:

neural_compressor/experimental/mixed_precision.py

@@ -131,6 +131,10 @@ def __call__(self):
             self._model.input_tensor_names = cfg.model.inputs if \
                 not self._input else self._input
             self._model.workspace_path = cfg.tuning.workspace.path
+            if 'bf16' in self._precisions or \
+               (cfg.mixed_precision and 'bf16' in cfg.mixed_precision.precisions) or \
+               (cfg.graph_optimization and 'bf16' in cfg.graph_optimization.precisions):
+                os.environ['MIX_PRECISION_TEST'] = '1'
 
         # when eval_func is set, will be directly used and eval_dataloader can be None
         if self._eval_func is None:

test/graph_optimization/test_graph_optimization.py

@@ -177,14 +177,10 @@ class TestGraphOptimizationOnNonBF16Host(unittest.TestCase):
     @classmethod
     def setUpClass(self):
         build_fake_yaml()
-        if CpuInfo().bf16:
-            os.environ['MIX_PRECISION_TEST'] = '1'
 
     @classmethod
     def tearDownClass(self):
         os.remove('fake_yaml.yaml')
-        if CpuInfo().bf16:
-           del os.environ['MIX_PRECISION_TEST']
 
     @disable_random()
     @unittest.skipIf(tf.__version__ < "2.0", "does not support on 1.15up3")
@@ -239,7 +235,6 @@ class TestGraphOptimization(unittest.TestCase):
     @classmethod
     def setUpClass(self):
         os.environ['FORCE_BF16'] = '1'
-        os.environ['MIX_PRECISION_TEST'] = '1'
         build_fake_yaml()
         build_fake_yaml_2()
         build_fake_yaml_3()
@@ -250,7 +245,6 @@ def setUpClass(self):
     @classmethod
     def tearDownClass(self):
         del os.environ['FORCE_BF16']
-        del os.environ['MIX_PRECISION_TEST']
         os.remove('fake_yaml.yaml')
         os.remove('fake_yaml_2.yaml')
         os.remove('fake_yaml_3.yaml')
@@ -968,7 +962,6 @@ def test_graph_optimization_without_yaml_with_precisions(self):
     @disable_random()
     def test_graph_optimization_fp32_only_with_force_bf16(self):
         os.environ['FORCE_BF16'] = '1'
-        os.environ['MIX_PRECISION_TEST'] = '1'
         x = tf.compat.v1.placeholder(tf.float32, [1, 56, 56, 16], name="input")
         top_relu = tf.nn.relu(x)
         paddings = tf.constant([[0, 0], [1, 1], [1, 1], [0, 0]])

test/mixed_precision/test_mixed_precision.py

@@ -293,7 +293,6 @@ class TestMixedPrecision(unittest.TestCase):
     def setUpClass(self):
         os.environ['FORCE_FP16'] = '1'
         os.environ['FORCE_BF16'] = '1'
-        os.environ['MIX_PRECISION_TEST'] = '1' 
         self.onnx_model = build_matmul_model()
         self.matmul_dataset = MatmulDataset()
         self.tf_model = build_tf_graph()
@@ -304,7 +303,6 @@ def setUpClass(self):
     def tearDownClass(self):
         del os.environ['FORCE_FP16']
         del os.environ['FORCE_BF16']
-        del os.environ['MIX_PRECISION_TEST']
         shutil.rmtree("./saved", ignore_errors=True)
         shutil.rmtree("./nc_workspace", ignore_errors=True)
         os.remove("test.yaml")