Quantization: Several issues with quantize_model function

Describe the bug

Hi, I’m aware of WIP on this project. I had several issues with quantize_model function. See my code examples below.

System information

TensorFlow installed from (source or binary): binary

TensorFlow version: 2.2.0-dev20200310

I use CPU version:

tf-estimator-nightly          2.1.0.dev2020031001 
tf-nightly-cpu                2.2.0.dev20200310   

TensorFlow Model Optimization version: 0.3.0.dev3

Python version: 3.7.0

Describe the expected behavior

I can quantize_model on any model. I should be able to reuse existing feature extractors.

Describe the current behavior

Examples are in section below.

Code to reproduce the issue Provide a reproducible code that is the bare minimum necessary to generate the problem.

I will show what work and what not.

1. This case work for me (I see QuantizeLayers in the summary):

import tensorflow_model_optimization as tfmot
import tensorflow as tf
keras = tf.keras

quantize_model = tfmot.quantization.keras.quantize_model
image_dim = 224

backbone = tf.keras.applications.MobileNetV2(
    input_shape=[image_dim, image_dim, 3],
    include_top=False,
    weights=None,
    input_tensor=None,
    pooling=None,
    classes=None,
)
q_model = quantize_model(backbone)

1. This one raises error:

input_image = keras.Input(shape=[image_dim, image_dim, 3], name='image')
outputs = backbone(input_image)
not_working_model = keras.Model(input_image, outputs)
q_model = quantize_model(not_working_model)

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-10-ff8232533be0> in <module>
      3 outputs = backbone(input_image)
      4 not_working_model = keras.Model(input_image, outputs)
----> 5 q_aware_model = quantize_model(not_working_model)

~/anaconda3/envs/tensorflow2/lib/python3.7/site-packages/tensorflow_model_optimization/python/core/quantization/keras/quantize.py in quantize_model(to_quantize)
     99         'the `quantize_annotate_layer` API to handle individual layers.')
    100 
--> 101   annotated_model = quantize_annotate_model(to_quantize)
    102   return quantize_apply(annotated_model)
    103 

~/anaconda3/envs/tensorflow2/lib/python3.7/site-packages/tensorflow_model_optimization/python/core/quantization/keras/quantize.py in quantize_annotate_model(to_annotate)
    148 
    149   return keras.models.clone_model(
--> 150       to_annotate, input_tensors=None, clone_function=_add_quant_wrapper)
    151 
    152 

~/anaconda3/envs/tensorflow2/lib/python3.7/site-packages/tensorflow/python/keras/models.py in clone_model(model, input_tensors, clone_function)
    425   else:
    426     return _clone_functional_model(
--> 427         model, input_tensors=input_tensors, layer_fn=clone_function)
    428 
    429 

~/anaconda3/envs/tensorflow2/lib/python3.7/site-packages/tensorflow/python/keras/models.py in _clone_functional_model(model, input_tensors, layer_fn)
    198   input_tensors, output_tensors, created_layers = (
    199       network.reconstruct_from_config(model_config,
--> 200                                       created_layers=created_layers))
    201   metrics_names = model.metrics_names
    202   model = Model(input_tensors, output_tensors, name=model.name)

~/anaconda3/envs/tensorflow2/lib/python3.7/site-packages/tensorflow/python/keras/engine/network.py in reconstruct_from_config(config, custom_objects, created_layers)
   2044     layer = created_layers[layer_name]
   2045     node_index = get_node_index(layer, node_index)
-> 2046     layer_output_tensors = layer._inbound_nodes[node_index].output_tensors
   2047     output_tensors.append(nest.flatten(layer_output_tensors)[tensor_index])
   2048 

TypeError: list indices must be integers or slices, not NoneType

1. This one seems to stuck in some infinite loop (I had to interrupt it):

input_image = keras.Input(shape=[image_dim, image_dim, 3], name='image')
outputs = backbone(input_image)
# add extra convolution
outputs = keras.layers.Conv2D(64, 3, activation='softmax')(outputs)
not_working_model = keras.Model(input_image, outputs)

# this seems to stuck in some while loop or something else
q_aware_model = quantize_model(not_working_model)

1. Creating custom keras model does not work:

input_image = keras.Input(shape=[image_dim, image_dim, 3], name='image')

class MyWrapper(tf.keras.Model):
    def call(self, x, **kwargs):
        return backbone(x) #  or return just x

outputs = MyWrapper()(input_image)
not_working_model = keras.Model(input_image, outputs)
q_aware_model = quantize_model(not_working_model)

Tail of traceback:

~/anaconda3/envs/tensorflow2/lib/python3.7/site-packages/tensorflow/python/keras/engine/network.py in get_network_config(network, serialize_layer_fn)
   2112           filtered_inbound_nodes.append(node_data)
   2113 
-> 2114     layer_config = serialize_layer_fn(layer)
   2115     layer_config['name'] = layer.name
   2116     layer_config['inbound_nodes'] = filtered_inbound_nodes

~/anaconda3/envs/tensorflow2/lib/python3.7/site-packages/tensorflow/python/keras/models.py in _copy_layer(layer)
    241       created_layers[layer.name] = InputLayer(**layer.get_config())
    242     else:
--> 243       created_layers[layer.name] = layer_fn(layer)
    244     return {}
    245 

~/anaconda3/envs/tensorflow2/lib/python3.7/site-packages/tensorflow/python/keras/models.py in _clone_layer(layer)
     59 
     60 def _clone_layer(layer):
---> 61   return layer.__class__.from_config(layer.get_config())
     62 
     63 

~/anaconda3/envs/tensorflow2/lib/python3.7/site-packages/tensorflow_model_optimization/python/core/quantization/keras/quantize_annotate.py in get_config(self)
     81 
     82   def get_config(self):
---> 83     base_config = super(QuantizeAnnotate, self).get_config()
     84     config = {'quantize_config': serialize_keras_object(self.quantize_config)}
     85     return dict(list(base_config.items()) + list(config.items()))

~/anaconda3/envs/tensorflow2/lib/python3.7/site-packages/tensorflow/python/keras/layers/wrappers.py in get_config(self)
     71         'layer': {
     72             'class_name': self.layer.__class__.__name__,
---> 73             'config': self.layer.get_config()
     74         }
     75     }

~/anaconda3/envs/tensorflow2/lib/python3.7/site-packages/tensorflow/python/keras/engine/network.py in get_config(self)
    960   def get_config(self):
    961     if not self._is_graph_network:
--> 962       raise NotImplementedError
    963     return copy.deepcopy(get_network_config(self))
    964 

NotImplementedError: 

1. Custom layer does not work for me:

input_image = keras.Input(shape=[image_dim, image_dim, 3], name='image')

class MyLayer(tf.keras.layers.Layer):
    def call(self, x, **kwargs):
        return x

outputs = MyLayer()(input_image)
not_working_model = keras.Model(input_image, outputs)
q_aware_model = quantize_model(not_working_model)

results in ValueError

~/anaconda3/envs/tensorflow2/lib/python3.7/site-packages/tensorflow/python/keras/utils/generic_utils.py in class_and_config_for_serialized_keras_object(config, module_objects, custom_objects, printable_module_name)
    319   cls = get_registered_object(class_name, custom_objects, module_objects)
    320   if cls is None:
--> 321     raise ValueError('Unknown ' + printable_module_name + ': ' + class_name)
    322 
    323   cls_config = config['config']

ValueError: Unknown layer: MyLayer

[alanchiao]

@Craftsman381, See https://www.tensorflow.org/model_optimization/guide/quantization/training_comprehensive_guide#checkpoint_and_deserialize for how to resolve your issue.

I had also submitted https://github.com/tensorflow/model-optimization/pull/344, so it’ll be clearer on what to do starting with the next TFMOT release.

[tensorbuffer]

@alanchiao Can you elaborate what are the “fundamental reasons we cannot support them well”? Do you have a timeline when QAT of subclassed model can be supported?