Mismatch in number of weights when loading quantized model (activation layer) #524
Description
Describe the bug
Saving and subsequently loading a quantized model results in the following error:
Traceback (most recent call last):
File "test.py", line 37, in <module>
model = tf.keras.models.load_model('MinimalExample.h5')
File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/saving/save.py", line 182, in load_model
return hdf5_format.load_model_from_hdf5(filepath, custom_objects, compile)
File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/saving/hdf5_format.py", line 181, in load_model_from_hdf5
load_weights_from_hdf5_group(f['model_weights'], model.layers)
File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/saving/hdf5_format.py", line 706, in load_weights_from_hdf5_group
str(len(weight_values)) + ' elements.')
ValueError: Layer #121 (named "quant_y_a_relu" in the current model) was found to correspond to layer quant_y_a_relu in the save file. However the new layer quant_y_a_relu expects 3 weights, but the saved weights have 1 elements.
The error can be reproduced using this code (test.py):
Please note that there is no error when setting quantize_model = False
import tensorflow as tf
import tensorflow_model_optimization as tfmot
from tensorflow_model_optimization.python.core.quantization.keras.default_8bit import default_8bit_quantize_configs
quantize_model = True
# Build model
base_model = tf.keras.applications.MobileNetV2(input_shape=(480,640,3),classes=2,weights='imagenet',include_top=False)
x = base_model.get_layer('block_12_add').output
y_a = tf.keras.layers.Conv2D(256,1,padding='same',dilation_rate=1,use_bias=False,kernel_initializer='he_normal',name='y_a_conv2d')(x)
y_a = tf.keras.layers.BatchNormalization(name='y_a_bn')(y_a)
y_a = tf.keras.layers.Activation('relu',name='y_a_relu')(y_a)
y_b = tf.keras.layers.Conv2D(256,3,padding='same',dilation_rate=6,use_bias=False,kernel_initializer='he_normal',name='y_b_conv2d')(x)
y_b = tf.keras.layers.BatchNormalization(name='y_b_bn')(y_b)
y_b = tf.keras.layers.Activation('relu',name='y_b_relu')(y_b)
output_tensor = tf.keras.layers.Concatenate(name='aspp_concat')([y_a,y_b])
model = tf.keras.models.Model(inputs=base_model.input,outputs=output_tensor,name='MinimalExample')
# Save model
optimizer = tf.keras.optimizers.Adam(learning_rate=1e-5)
loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
metrics = ['accuracy']
if quantize_model:
with tf.keras.utils.custom_object_scope({'NoOpQuantizeConfig':default_8bit_quantize_configs.NoOpQuantizeConfig}):
model = tfmot.quantization.keras.quantize_model(model)
model.compile(optimizer=optimizer,loss=loss,metrics=metrics)
model.save('MinimalExample.h5')
del model
# Load model
if quantize_model:
with tfmot.quantization.keras.quantize_scope():
model = tf.keras.models.load_model('MinimalExample.h5')
else:
model = tf.keras.models.load_model('MinimalExample.h5')
# Convert model
converter = tf.lite.TFLiteConverter.from_keras_model(model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
tflite_model = converter.convert()
# Export model
open('MinimalExample.tflite','wb').write(tflite_model)
System information
TensorFlow version (installed from source or binary): 2.3.0 (Docker image)
TensorFlow Model Optimization version (installed from source or binary): 0.4.1 (Docker image)
Python version: 3.6.9 (Docker image)
Describe the expected behavior
Code should not crash
Describe the current behavior
Code crashes