Add better error messages when attempting to quantize Lambda layers.

PiperOrigin-RevId: 415179482

Author: daverim

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

@@ -13,6 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 """Quantization API functions for tf.keras models."""
+import warnings
 
 import tensorflow as tf
 
@@ -177,6 +178,9 @@ def quantize_annotate_model(to_annotate):
     New tf.keras model with each layer in the model wrapped with
     `QuantizeAnnotate`. The new model preserves weights from the original
     model.
+
+  Raises:
+    ValueError: if the model cannot be annotated.
   """
   if to_annotate is None:
     raise ValueError('`to_annotate` cannot be None')
@@ -200,6 +204,15 @@ def _add_quant_wrapper(layer):
     if isinstance(layer, quantize_annotate_mod.QuantizeAnnotate):
       return layer
 
+    if isinstance(layer, tf.keras.layers.Lambda):
+      warnings.warn(
+          'Lambda layers are not supported by automatic model annotation '
+          'because the internal functionality cannot always be determined by '
+          'serialization alone. We recommend that you make a custom layer '
+          'and add a custom QuantizeConfig for it instead. This layer will not '
+          'be quantized which may lead to unexpected results.')
+      return layer
+
     if isinstance(layer, tf.keras.Model):
       raise ValueError(
           'Quantizing a tf.keras Model inside another tf.keras Model is not supported.'

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

@@ -187,6 +187,21 @@ def testQuantizeAnnotateModel_FailsWithNestedModels(self):
           keras.Sequential(
               [keras.Sequential([keras.layers.Dense(10, input_shape=(2,))])]))
 
+  def testQuantizeAnnotateModel_SkipsLambda(self):
+    model = keras.Sequential([
+        keras.layers.Dense(10, input_shape=(5,)),
+        keras.layers.Dropout(0.4),
+        keras.layers.Lambda(lambda x: x + 1.0)
+    ])
+    with self.assertWarns(Warning):
+      annotated_model = quantize_annotate_model(model)
+
+    self._assertWrappedLayer(annotated_model.layers[0])
+    self._assertWrappedLayer(annotated_model.layers[1])
+    self.assertIsInstance(annotated_model.layers[2], keras.layers.Lambda)
+    inputs = np.random.rand(1, 5)
+    self.assertAllEqual(model.predict(inputs), annotated_model.predict(inputs))
+
 
 class QuantizeApplyTest(tf.test.TestCase):