Support quantization of gelu activation

PiperOrigin-RevId: 424128991

Author: daverim

tensorflow_model_optimization/python/core/quantization/keras/default_8bit/default_8bit_quantize_registry.py

@@ -491,7 +491,7 @@ def get_output_quantizers(self, layer):
                        'Default8BitActivationQuantizeConfig.'.format(
                            layer.activation))
 
-    if layer.activation.__name__ in ['relu', 'swish']:
+    if layer.activation.__name__ in ['relu', 'swish', 'gelu']:
       # 'relu' should generally get fused into the previous layer.
       return [quantizers.MovingAverageQuantizer(
           num_bits=8, per_axis=False, symmetric=False, narrow_range=False)]

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

@@ -75,12 +75,13 @@ class QuantizeAwareActivation(object):
 
   # TODO(pulkitb): Other activations such as elu, tanh etc., should just work
   # on inclusion. Verify in TFLite before enabling.
+  # gelu requires both because it is not folded by tflite.
 
   # These activations should be quantized prior to the activation being applied.
-  _PRE_QUANT_ACTIVATIONS = frozenset({'softmax', 'sigmoid', 'tanh'})
+  _PRE_QUANT_ACTIVATIONS = frozenset({'softmax', 'sigmoid', 'tanh', 'gelu'})
 
   # These activations should be quantized after the activation has been applied.
-  _POST_QUANT_ACTIVATIONS = frozenset({'linear', 'relu', 'swish'})
+  _POST_QUANT_ACTIVATIONS = frozenset({'linear', 'relu', 'swish', 'gelu'})
 
   # Don't take any quantize operations for these activations.
   _NO_QUANTIZE_ACTIVATIONS = frozenset({'NoOpActivation'})

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

@@ -59,6 +59,15 @@ def call(self, inputs, training=None):
     def compute_output_shape(self, input_shape):
       return input_shape
 
+  def testSupportedPreAndPostActivation(self):
+    layer = self.TestLayer()
+    layer.activation = QuantizeAwareActivation(
+        activations.get('gelu'), self.quantizer, 0, layer)
+    model = keras.Sequential([layer])
+    names = ', '.join([weight.name for weight in model.layers[-1].weights])
+    self.assertIn('pre_activation', names)
+    self.assertIn('post_activation', names)
+
   def testConstruction_SupportedAndUnsupportedActivations(self):
     layer = self.TestLayer()