Disable all quantize layers where we don't have a tested conversion path to full-integer deployment, assuming TFLite enables uint8 inputs and outputs.

Will reenable once path is confirmed.

PiperOrigin-RevId: 304478647

Author: alanchiao

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

@@ -72,6 +72,8 @@ def _get_rnn_cells(self, rnn_layer):
 class QuantizeRegistry(quantize_registry.QuantizeRegistry, _RNNHelper):
   """QuantizationRegistry for built-in Keras classes for default 8-bit scheme."""
 
+  # TODO(tfmot): expand layers test in quantize_functional_test.py
+  # to add more layers to whitelist.
   _LAYER_QUANTIZE_INFO = [
 
       # Activation Layers
@@ -84,16 +86,18 @@ class QuantizeRegistry(quantize_registry.QuantizeRegistry, _RNNHelper):
       # layers.ThresholdedReLU,
 
       # Convolution Layers
-      _QuantizeInfo(layers.Conv1D, ['kernel'], ['activation']),
-      _QuantizeInfo(layers.Conv3D, ['kernel'], ['activation']),
-      # TODO(pulkitb): Verify Transpose layers.
-      _QuantizeInfo(layers.Conv2DTranspose, ['kernel'], ['activation']),
-      _QuantizeInfo(layers.Conv3DTranspose, ['kernel'], ['activation']),
+      # _QuantizeInfo(layers.Conv1D, ['kernel'], ['activation']),
+
+      # layers.Conv2D is supported and handled in code below.
+
+      # _QuantizeInfo(layers.Conv3D, ['kernel'], ['activation']),
+      # _QuantizeInfo(layers.Conv2DTranspose, ['kernel'], ['activation']),
+      # _QuantizeInfo(layers.Conv3DTranspose, ['kernel'], ['activation']),
       _no_quantize(layers.Cropping1D),
       _no_quantize(layers.Cropping2D),
       _no_quantize(layers.Cropping3D),
       _no_quantize(layers.UpSampling1D),
-      _no_quantize(layers.UpSampling2D),
+      # _no_quantize(layers.UpSampling2D),
       _no_quantize(layers.UpSampling3D),
       _no_quantize(layers.ZeroPadding1D),
       _no_quantize(layers.ZeroPadding2D),
@@ -107,7 +111,7 @@ class QuantizeRegistry(quantize_registry.QuantizeRegistry, _RNNHelper):
       _QuantizeInfo(layers.Dense, ['kernel'], ['activation']),
       _no_quantize(layers.Dropout),
       _no_quantize(layers.Flatten),
-      _no_quantize(layers.Masking),
+      # _no_quantize(layers.Masking),
       _no_quantize(layers.Permute),
       _no_quantize(layers.RepeatVector),
       _no_quantize(layers.Reshape),
@@ -119,7 +123,7 @@ class QuantizeRegistry(quantize_registry.QuantizeRegistry, _RNNHelper):
       # Pooling Layers
       _QuantizeInfo(layers.AveragePooling1D, [], [], True),
       _QuantizeInfo(layers.AveragePooling2D, [], [], True),
-      _QuantizeInfo(layers.AveragePooling3D, [], [], True),
+      # _QuantizeInfo(layers.AveragePooling3D, [], [], True),
       _QuantizeInfo(layers.GlobalAveragePooling1D, [], [], True),
       _QuantizeInfo(layers.GlobalAveragePooling2D, [], [], True),
       _QuantizeInfo(layers.GlobalAveragePooling3D, [], [], True),
@@ -128,11 +132,10 @@ class QuantizeRegistry(quantize_registry.QuantizeRegistry, _RNNHelper):
       _no_quantize(layers.GlobalMaxPooling3D),
       _no_quantize(layers.MaxPooling1D),
       _no_quantize(layers.MaxPooling2D),
-      _no_quantize(layers.MaxPooling3D),
+      # _no_quantize(layers.MaxPooling3D),
 
-      # TODO(pulkitb): Verify Locally Connected layers.
-      _QuantizeInfo(layers.LocallyConnected1D, ['kernel'], ['activation']),
-      _QuantizeInfo(layers.LocallyConnected2D, ['kernel'], ['activation']),
+      # _QuantizeInfo(layers.LocallyConnected1D, ['kernel'], ['activation']),
+      # _QuantizeInfo(layers.LocallyConnected2D, ['kernel'], ['activation']),
       _QuantizeInfo(layers.Add, [], [], True),
 
       # Enable once verified with TFLite behavior.

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

@@ -68,41 +68,43 @@ def testQuantizesWeightsInLayer(self):
   #  1. Only layers with 'kernel' attribute work. Need to extend to others.
   #  2. Activations are not tested currently.
   #  3. RNN layers need to be added
-
+  # TODO(tfmot): reenable some of these once added back to whitelist
+  # after testing in quantize_functional_test.py.
   @parameterized.parameters(
-      (layers.Conv1D, (3, 6), {
-          'filters': 4,
-          'kernel_size': 2
-      }),
+      # (layers.Conv1D, (3, 6), {
+      #     'filters': 4,
+      #     'kernel_size': 2
+      # }),
       (layers.Conv2D, (4, 6, 1), {
           'filters': 4,
           'kernel_size': (2, 2)
       }),
-      (layers.Conv2DTranspose, (7, 6, 3), {
-          'filters': 2,
-          'kernel_size': (3, 3)
-      }),
-      (layers.Conv3D, (5, 7, 6, 3), {
-          'filters': 2,
-          'kernel_size': (3, 3, 3)
-      }),
-      (layers.Conv3DTranspose, (5, 7, 6, 3), {
-          'filters': 2,
-          'kernel_size': (3, 3, 3)
-      }),
+      # (layers.Conv2DTranspose, (7, 6, 3), {
+      #     'filters': 2,
+      #     'kernel_size': (3, 3)
+      # }),
+      # (layers.Conv3D, (5, 7, 6, 3), {
+      #     'filters': 2,
+      #     'kernel_size': (3, 3, 3)
+      # }),
+      # (layers.Conv3DTranspose, (5, 7, 6, 3), {
+      #     'filters': 2,
+      #    'kernel_size': (3, 3, 3)
+      # }),
       # TODO(pulkitb): Add missing SeparableConv layers. The initializers are
       # different, so will need a change.
       (layers.Dense, (3,), {
           'units': 2
       }),
-      (layers.LocallyConnected1D, (3, 6), {
-          'filters': 4,
-          'kernel_size': 2
-      }),
-      (layers.LocallyConnected2D, (4, 6, 1), {
-          'filters': 4,
-          'kernel_size': (2, 2)
-      }))
+      # (layers.LocallyConnected1D, (3, 6), {
+      #     'filters': 4,
+      #     'kernel_size': 2
+      # }),
+      # (layers.LocallyConnected2D, (4, 6, 1), {
+      #     'filters': 4,
+      #     'kernel_size': (2, 2)
+      # })
+  )
   def testQuantizesWeights_KerasLayers(self, layer_type, input_shape, kwargs):
     self.weights = None