Add sigmoid as supported activation in QAT

Handling of sigmoid is similar to Softmax. We place a FQ before
the activation but not after to prevent a large set of values turning
to zeros, and potentially leading to NaNs.

PiperOrigin-RevId: 336213731

Author: nutsiepully

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

@@ -472,7 +472,7 @@ def get_output_quantizers(self, layer):
       # 'relu' should generally get fused into the previous layer.
       return [quantizers.MovingAverageQuantizer(
           num_bits=8, per_axis=False, symmetric=False, narrow_range=False)]
-    elif layer.activation.__name__ in ['linear', 'softmax']:
+    elif layer.activation.__name__ in ['linear', 'softmax', 'sigmoid']:
       return []
 
     raise ValueError('Activation {} not supported by '

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

@@ -77,7 +77,7 @@ class QuantizeAwareActivation(object):
   # on inclusion. Verify in TFLite before enabling.
 
   # These activations should be quantized prior to the activation being applied.
-  _PRE_QUANT_ACTIVATIONS = frozenset({'softmax'})
+  _PRE_QUANT_ACTIVATIONS = frozenset({'softmax', 'sigmoid'})
 
   # These activations should be quantized after the activation has been applied.
   _POST_QUANT_ACTIVATIONS = frozenset({'linear', 'relu'})