Improve or add error messages for nested models and subclassed models.

Context: #292.
PiperOrigin-RevId: 303861624

Author: alanchiao

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

@@ -104,6 +104,12 @@ def quantize_model(to_quantize):
         'You passed an instance of type: {input}.'.format(
             input=to_quantize.__class__.__name__))
 
+  if not isinstance(to_quantize, keras.Sequential) \
+      and not to_quantize._is_graph_network:  # pylint: disable=protected-access
+    raise ValueError(
+        '`to_quantize` can only either be a tf.keras Sequential or '
+        'Functional model.')
+
   annotated_model = quantize_annotate_model(to_quantize)
   return quantize_apply(annotated_model)
 
@@ -149,11 +155,23 @@ def quantize_annotate_model(to_annotate):
         'You passed an instance of type: {input}.'.format(
             input=to_annotate.__class__.__name__))
 
+  if not isinstance(to_annotate, keras.Sequential) \
+      and not to_annotate._is_graph_network:  # pylint: disable=protected-access
+    raise ValueError(
+        '`to_annotate` can only either be a tf.keras Sequential or '
+        'Functional model.')
+
   def _add_quant_wrapper(layer):
+    """Add annotation wrapper."""
     # Already annotated layer. No need to wrap.
     if isinstance(layer, quantize_annotate_mod.QuantizeAnnotate):
       return layer
 
+    if isinstance(layer, tf.keras.Model):
+      raise ValueError(
+          'Quantizing a tf.keras Model inside another tf.keras Model is not supported.'
+      )
+
     return quantize_annotate_mod.QuantizeAnnotate(layer)
 
   return keras.models.clone_model(

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

@@ -52,6 +52,17 @@ def __init__(self, layer, quantize_config=None, **kwargs):
     """
     super(QuantizeAnnotate, self).__init__(layer, **kwargs)
 
+    if layer is None:
+      raise ValueError('`layer` cannot be None.')
+
+    # Check against keras.Model since it is an instance of keras.layers.Layer.
+    if not isinstance(layer, tf.keras.layers.Layer) or isinstance(
+        layer, tf.keras.Model):
+      raise ValueError(
+          '`layer` can only be a `tf.keras.layers.Layer` instance. '
+          'You passed an instance of type: {input}.'.format(
+              input=layer.__class__.__name__))
+
     self.quantize_config = quantize_config
 
     self._track_trackable(layer, name='layer')

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

@@ -89,6 +89,12 @@ def testSerializationQuantizeAnnotate(self):
 
     self.assertEqual(wrapper_from_config.get_config(), wrapper.get_config())
 
+  def testQuantizeAnnotate_FailsWithModel(self):
+    layer = keras.layers.Dense(5, activation='relu', input_shape=(10,))
+    model = keras.Sequential([layer])
+
+    with self.assertRaises(ValueError):
+      quantize_annotate.QuantizeAnnotate(model)
 
 if __name__ == '__main__':
   tf.test.main()

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

@@ -165,6 +165,20 @@ def testQuantizeAnnotateModel_RemovesOptimizer(self):
     annotated_model = quantize_annotate_model(model)
     self.assertIsNone(annotated_model.optimizer)
 
+  def testQuantizeAnnotateModel_FailsWithSubclassedModel(self):
+    class MyModel(keras.Model):
+      def call(self, inputs, training=None, mask=None):  # pylint: disable=g-wrong-blank-lines
+        return inputs
+
+    with self.assertRaises(ValueError):
+      quantize_annotate_model(MyModel())
+
+  def testQuantizeAnnotateModel_FailsWithNestedModels(self):
+    with self.assertRaises(ValueError):
+      quantize_annotate_model(
+          keras.Sequential(
+              [keras.Sequential([keras.layers.Dense(10, input_shape=(2,))])]))
+
 
 class QuantizeApplyTest(tf.test.TestCase):
 
@@ -411,6 +425,14 @@ def testQuantizeApply_RemovesOptimizer(self):
     quantized_model = quantize_apply(annotated_model)
     self.assertIsNone(quantized_model.optimizer)
 
+  def testQuantizeApply_RunsWhenNestedModelNotAnnotated(self):
+    annotated_model = keras.Sequential([
+        keras.Sequential([keras.layers.Dense(10, input_shape=(2,))]),
+        quantize_annotate_layer(keras.layers.Dense(10)),
+    ])
+
+    quantize_apply(annotated_model)
+
 
 if __name__ == '__main__':
   tf.test.main()

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

@@ -47,6 +47,16 @@ def __init__(self, layer, quantize_config, **kwargs):
       quantize_config: `QuantizeConfig` to quantize layer.
       **kwargs: Additional keyword arguments to be passed to the keras layer.
     """
+    if layer is None:
+      raise ValueError('`layer` cannot be None.')
+
+    # Check against keras.Model since it is an instance of keras.layers.Layer.
+    if not isinstance(layer, tf.keras.layers.Layer) or isinstance(
+        layer, tf.keras.Model):
+      raise ValueError(
+          '`layer` can only be a `tf.keras.layers.Layer` instance. '
+          'You passed an instance of type: {input}.'.format(
+              input=layer.__class__.__name__))
 
     if quantize_config is None:
       raise ValueError('quantize_config cannot be None. It is needed to '

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

@@ -180,6 +180,15 @@ def testSerializationQuantizeWrapper(self):
 
     self.assertEqual(wrapper_from_config.get_config(), wrapper.get_config())
 
+  def testQuantizeWrapper_FailsWithModel(self):
+    layer = keras.layers.Dense(5, activation='relu', input_shape=(10,))
+    model = keras.Sequential([layer])
+
+    with self.assertRaises(ValueError):
+      QuantizeWrapper(
+          model,
+          quantize_config=self.quantize_registry.get_quantize_config(layer))
+
   # TODO(pulkitb): Add test to ensure weights are also preserved.