Add operations before Identity operations should be quantized.

Fixes #19014

PiperOrigin-RevId: 195443326

2 files changed, 14 insertions, 12 deletions

diff --git a/tensorflow/contrib/quantize/python/quantize.py b/tensorflow/contrib/quantize/python/quantize.py
index efc1a94b3c..60616ea749 100644
--- a/tensorflow/contrib/quantize/python/quantize.py
+++ b/tensorflow/contrib/quantize/python/quantize.py
@@ -33,7 +33,7 @@ from tensorflow.python.platform import tf_logging as logging
 _QUANTIZABLE_TYPES = {'Conv2D', 'MatMul', 'DepthwiseConv2dNative'}
 
 # Activations that are supported by the quantization rewrite.
-_ACTIVATION_TYPES = {'Relu', 'Relu6', 'Identity'}
+_ACTIVATION_TYPES = {'Relu', 'Relu6'}
 
 
 def Quantize(graph,
@@ -267,8 +267,10 @@ def _FindLayersToQuantize(graph):
 
   # The input to the activation can come from bias add, fold bias add, the
   # bypasses.
+  # TODO(suharshs): We should ideally skip Identity operations instead of
+  # treating them as an activation.
   activation_pattern = graph_matcher.OpTypePattern(
-      '|'.join(_ACTIVATION_TYPES),
+      '|'.join(_ACTIVATION_TYPES) + '|Identity',
       inputs=[
           graph_matcher.OneofPattern([
               bias_add_pattern, folded_bias_add_pattern, bypass_pattern_a,
diff --git a/tensorflow/contrib/quantize/python/quantize_test.py b/tensorflow/contrib/quantize/python/quantize_test.py
index 5e479f3946..e7360ae03c 100644
--- a/tensorflow/contrib/quantize/python/quantize_test.py
+++ b/tensorflow/contrib/quantize/python/quantize_test.py
@@ -74,7 +74,7 @@ class QuantizeTest(test_util.TensorFlowTestCase):
                     weights_initializer=self._WeightInit(0.09),
                     activation_fn=None, scope='test/test')
       node = math_ops.add(conv, input2, name='test/add')
-      node = array_ops.identity(node, name='test/identity')
+      node = nn_ops.relu6(node, name='test/relu6')
       update_barrier = control_flow_ops.no_op(name='update_barrier')
       with ops.control_dependencies([update_barrier]):
         array_ops.identity(node, name='control_dependency')
@@ -97,7 +97,7 @@ class QuantizeTest(test_util.TensorFlowTestCase):
         for output in quant_op.outputs:
           consumers.extend(output.consumers())
 
-        self.assertNotIn('test/identity', [c.name for c in consumers])
+        self.assertNotIn('test/relu6', [c.name for c in consumers])
 
   def testInsertQuantOpForAddAfterSeparableConv2d(self):
     self._RunTestOverParameters(
@@ -114,7 +114,7 @@ class QuantizeTest(test_util.TensorFlowTestCase):
                               weights_initializer=self._WeightInit(0.09),
                               activation_fn=None, scope='test/test')
       node = math_ops.add(conv, input2, name='test/add')
-      node = array_ops.identity(node, name='test/identity')
+      node = nn_ops.relu6(node, name='test/relu6')
       update_barrier = control_flow_ops.no_op(name='update_barrier')
       with ops.control_dependencies([update_barrier]):
         array_ops.identity(node, name='control_dependency')
@@ -135,7 +135,7 @@ class QuantizeTest(test_util.TensorFlowTestCase):
         for output in quant_op.outputs:
           consumers.extend(output.consumers())
 
-        self.assertNotIn('test/identity', [c.name for c in consumers])
+        self.assertNotIn('test/relu6', [c.name for c in consumers])
 
   def testFinalLayerQuantized(self):
     self._RunTestOverParameters(self._TestFinalLayerQuantized)
@@ -174,7 +174,7 @@ class QuantizeTest(test_util.TensorFlowTestCase):
           stride=2,
           padding='SAME',
           weights_initializer=self._WeightInit(0.09),
-          activation_fn=array_ops.identity,
+          activation_fn=nn_ops.relu6,
           scope='test/test')
       bypass_tensor = math_ops.add(conv, input2, name='test/add')
       # The output of the post_activation bypass will be another layer.
@@ -184,7 +184,7 @@ class QuantizeTest(test_util.TensorFlowTestCase):
           stride=2,
           padding='SAME',
           weights_initializer=self._WeightInit(0.09),
-          activation_fn=array_ops.identity,
+          activation_fn=nn_ops.relu6,
           scope='test/unused')
 
       quantize.Quantize(graph, is_training, weight_bits=8, activation_bits=8)
@@ -212,7 +212,7 @@ class QuantizeTest(test_util.TensorFlowTestCase):
           stride=2,
           padding='SAME',
           weights_initializer=self._WeightInit(0.09),
-          activation_fn=array_ops.identity,
+          activation_fn=nn_ops.relu6,
           scope='test/test1')
 
       # The bypass of this conv is the post activation bypass of the previous
@@ -227,7 +227,7 @@ class QuantizeTest(test_util.TensorFlowTestCase):
           scope='test/test2')
 
       bypass_tensor = math_ops.add(conv1, conv2, name='test/add')
-      _ = array_ops.identity(bypass_tensor, name='test/output')
+      _ = nn_ops.relu6(bypass_tensor, name='test/output')
 
       quantize.Quantize(graph, is_training, weight_bits=8, activation_bits=8)
 
@@ -248,11 +248,11 @@ class QuantizeTest(test_util.TensorFlowTestCase):
           'test/test1/act_quant/FakeQuantWithMinMaxVars' in op_names)
       self.assertTrue('test/act_quant/FakeQuantWithMinMaxVars' in op_names)
       self.assertEqual(
-          'Identity',
+          'Relu6',
           graph.get_operation_by_name(
               'test/test1/act_quant/FakeQuantWithMinMaxVars').inputs[0].op.type)
       self.assertEqual(
-          'Identity',
+          'Relu6',
           graph.get_operation_by_name(
               'test/act_quant/FakeQuantWithMinMaxVars').inputs[0].op.type)