Removing activation quantize code.

Code to apply transformations will now go into a separate
interface which will be added in future CLs. Wrapping using
QuantizeAwareActivation will be done in the QuantizeWrapper.

Removing this code, before adding the features back in in the
correct place.

PiperOrigin-RevId: 263803284

Author: nutsiepully

tensorflow_model_optimization/python/core/quantization/keras/quantize_emulate.py

@@ -168,24 +168,6 @@ def _clone_model_with_weights(model_to_clone):
 
     return cloned_model
 
-  def _quantize_activation(activation, parent_class, quantize_params):
-    try:
-      return quantize_aware_activation.QuantizeAwareActivation(
-          activation.__name__, parent_class, **quantize_params)
-    except TypeError:
-      # Non-standard activation. Could be a custom callable, or an advanced
-      # activation. Simply return the original activation for now.
-      # TODO(pulkitb): Determine how to handle custom activations and advanced
-      # activations.
-      return activation
-
-  def _get_quantize_activation_params(layer):
-    quant_params = layer.get_quantize_params()
-    # narrow_range is not relevant to quantizing activations.
-    quant_params.pop('narrow_range')
-
-    return quant_params
-
   def _apply_quantization(quant_annotate_layer):
     return QuantizeEmulateWrapper(
         quant_annotate_layer.layer,
@@ -195,32 +177,7 @@ def _apply_quantization(quant_annotate_layer):
   # model without modifying the weights of the original model.
   model_copy = _clone_model_with_weights(model)
 
-  # Apply all graph level transformations.
-  replace_map = {}
-
-  # Replace activations in layers with QuantAwareActivation.
-  # Dense(activation='relu') -> Dense(activation=QuantAwareActivation('relu'))
-  # TODO(pulkitb): Not all layers (LSTMs) have just activation. Add
-  # generic handling for all layers.
-  for layer in model_copy.layers:
-    if isinstance(layer, quant_annotate.QuantizeAnnotate) and \
-        (layer.layer.activation is not None and
-         layer.layer.activation != keras.activations.linear):
-      quantized_layer = _apply_quantization(layer)
-
-      quantized_layer.layer.activation = _quantize_activation(
-          layer.layer.activation, layer.layer.__class__,
-          _get_quantize_activation_params(layer))
-
-      replace_map[layer] = quantized_layer
-
-  # TODO(pulkitb): Transform [Dense(), ReLU()] to be quant aware.
-
   def _add_quant_emulate_wrapper(layer):  # pylint: disable=missing-docstring
-    # Quantized layer has been constructed during graph transformation. Return.
-    if layer in replace_map:
-      return replace_map[layer]
-
     if not isinstance(layer, quant_annotate.QuantizeAnnotate):
       return layer
 

tensorflow_model_optimization/python/core/quantization/keras/quantize_emulate_test.py

@@ -32,7 +32,6 @@
 quantize_annotate = quantize_emulate.quantize_annotate
 QuantizeEmulate = quantize_emulate.QuantizeEmulate
 QuantizeEmulateWrapper = quantize_emulate_wrapper.QuantizeEmulateWrapper
-QuantizeAwareActivation = quantize_aware_activation.QuantizeAwareActivation
 
 
 class QuantizeEmulateTest(test.TestCase):
@@ -270,48 +269,6 @@ def testAppliesQuantizationToAnnotatedModel_Functional(self):
 
     self._assert_model_emulated(model, quantized_model)
 
-  # Transformation Tests
-
-  def testQuantizesActivationsWithinLayer_Sequential(self):
-    quant_params = {'num_bits': 8, 'symmetric': True}
-    model = keras.Sequential([
-        quantize_annotate(
-            keras.layers.Conv2D(32, 5, activation='relu'),
-            input_shape=(28, 28, 1))
-    ])
-
-    quantized_model = quantize_emulate.quantize_apply(model)
-
-    # We expect activation to be modified.
-    self._assert_model_emulated(model, quantized_model, ['activation'])
-
-    conv_layer = quantized_model.layers[0].layer
-    self.assertIsInstance(conv_layer.activation, QuantizeAwareActivation)
-    self.assertEqual(
-        keras.activations.get('relu'), conv_layer.activation.activation)
-    self.assertEqual(keras.layers.Conv2D, conv_layer.activation.parent_layer)
-    self.assertEqual(quant_params, conv_layer.activation.get_quantize_params())
-
-  def testQuantizesActivationsWithinLayer_Functional(self):
-    quant_params = {'num_bits': 8, 'symmetric': True}
-
-    inputs = keras.Input(shape=(28, 28, 1))
-    results = quantize_annotate(
-        keras.layers.Conv2D(32, 5, activation='relu'))(inputs)
-    model = keras.Model(inputs=inputs, outputs=results)
-
-    quantized_model = quantize_emulate.quantize_apply(model)
-
-    # We expect activation to be modified.
-    self._assert_model_emulated(model, quantized_model, ['activation'])
-
-    conv_layer = quantized_model.layers[1].layer
-    self.assertIsInstance(conv_layer.activation, QuantizeAwareActivation)
-    self.assertEqual(
-        keras.activations.get('relu'), conv_layer.activation.activation)
-    self.assertEqual(keras.layers.Conv2D, conv_layer.activation.parent_layer)
-    self.assertEqual(quant_params, conv_layer.activation.get_quantize_params())
-
 
 if __name__ == '__main__':
   test.main()