ValueError：未知层：AnchorBoxes 量化 tensorflow

Question

I am applying quantization to a SSD model. The gist is attached. There is a custom object called "AnchorBoxes" which is added while loading the model. This works fine when I don't do quantization. But when I apply quantization, this custom object is not recognized.

I tried a work around.

def apply_quantization_to_conv2D(layer):
  #print(layer)
  if isinstance(layer, tf.keras.layers.Conv2D):
    return tfmot.quantization.keras.quantize_annotate_layer(layer)
  return layer

# Use `tf.keras.models.clone_model` to apply `apply_quantization_to_dense` 
# to the layers of the model.
annotated_model = tf.keras.models.clone_model(
    model,
    clone_function=apply_quantization_to_conv2D,
)

#annotated_model.save('quantize_ready_model_20_01_Conv2D.h5', include_optimizer=True)
annotated_model.summary()
# Now that the Dense layers are annotated,
# `quantize_apply` actually makes the model quantization aware.



#quant_aware_model = tfmot.quantization.keras.quantize_apply(annotated_model)

I commented this line quant_aware_model = tfmot.quantization.keras.quantize_apply(annotated_model) in the above code as it was throwing the error ValueError: Unknown layer: AnchorBoxes

Instead I saved the model after applying quantization to the Conv2D layers as below

def apply_quantization_to_conv2D(layer):
  #print(layer)
  if isinstance(layer, tf.keras.layers.Conv2D):
    return tfmot.quantization.keras.quantize_annotate_layer(layer)
  return layer

# Use `tf.keras.models.clone_model` to apply `apply_quantization_to_dense` 
# to the layers of the model.
annotated_model = tf.keras.models.clone_model(
    model,
    clone_function=apply_quantization_to_conv2D,
)


annotated_model.summary()

annotated_model.save('quantize_ready_model_20_01_Conv2D_1.h5', include_optimizer=True)
# Now that the Dense layers are annotated,
# `quantize_apply` actually makes the model quantization aware.



#quant_aware_model = tfmot.quantization.keras.quantize_apply(annotated_model)
#quant_aware_model.compile(optimizer=adam, loss=ssd_loss.compute_loss)
#quant_aware_model.summary()

Then I loaded the model hoping that the loaded quantized model as below will have the custom_objects attached to it.

with tfmot.quantization.keras.quantize_scope():
    loaded_model = tf.keras.models.load_model('./quantize_ready_model_20_01_Conv2D_1.h5', custom_objects={'AnchorBoxes': AnchorBoxes})

Finally I applied the quantize_apply to the new loaded_model which has quantized layers.

quant_aware_model = tfmot.quantization.keras.quantize_apply(loaded_model)

which again resulted in the same error

ValueError: Unknown layer: AnchorBoxes

System information

TensorFlow version (installed from source or binary): TF 2.0.0

TensorFlow Model Optimization version (installed from source or binary): 0.5.0

Describe the expected behavior
When I run quantize_apply(model), the model should become quantization aware

Describe the current behavior
Throwing an error on the custom objects

Code to reproduce the issue
gist

Answer 1

The issue was fixed after passing the custom layer like this AnchorBoxes': AnchorBoxes in the below code.

with quantize_scope(
  {'DefaultDenseQuantizeConfig': DefaultDenseQuantizeConfig,
   'AnchorBoxes': AnchorBoxes}):
  # Use `quantize_apply` to actually make the model quantization aware.
  quant_aware_model = tfmot.quantization.keras.quantize_apply(annotated_model)