use tensorflow.compat.v1

src/sdk/pynni/nni/compression/tensorflow/builtin_pruners.py

@@ -1,5 +1,5 @@
 import logging
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 from .compressor import Pruner
 
 __all__ = [ 'LevelPruner', 'AGP_Pruner', 'SensitivityPruner' ]

src/sdk/pynni/nni/compression/tensorflow/builtin_quantizers.py

@@ -1,5 +1,5 @@
 import logging
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 from .compressor import Quantizer
 
 __all__ = [ 'NaiveQuantizer', 'QAT_Quantizer', 'DoReFaQuantizer' ]

src/sdk/pynni/nni/compression/tensorflow/compressor.py

@@ -1,4 +1,4 @@
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 import logging
 from . import default_layers
 

src/sdk/pynni/tests/test_compressor.py

@@ -1,10 +1,12 @@
 from unittest import TestCase, main
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 import torch
 import torch.nn.functional as F
 import nni.compression.tensorflow as tf_compressor
 import nni.compression.torch as torch_compressor
 
+tf.disable_v2_behavior()
+
 def weight_variable(shape):
     return tf.Variable(tf.truncated_normal(shape, stddev = 0.1))