使用TensorFlow Hub模型进行量化感知训练时遇KerasLayer未知层错误

问题背景 #

使用TensorFlow Hub中的模型完成自定义数据集的迁移学习后，尝试进行量化感知训练时出现错误，相关代码及报错信息如下：

代码片段 #

import numpy as np
import time

import PIL.Image as Image
import matplotlib.pylab as plt

import tensorflow as tf
from tensorflow import keras
from keras.preprocessing.image import ImageDataGenerator

import tensorflow_hub as hub

import datetime


from google.colab import drive
drive.mount('/content/gdrive', force_remount=True)
dataset_path = '/content/gdrive/MyDrive/images/'

image_size = (224,224)
batch_size = 10

train_datagen = ImageDataGenerator(rescale=1./255, 
                                   rotation_range=15,
                                   zoom_range = (0.95,0.95),
                                   width_shift_range=0.1,
                                   height_shift_range=0.1,
                                   validation_split=0.2,
                                   dtype = tf.float32,
                                  )

validation_datagen = ImageDataGenerator(rescale=1./255, validation_split=0.2, dtype = tf.float32,)

train_batches = train_datagen.flow_from_directory(
    dataset_path,
    target_size = image_size,
    batch_size = batch_size,
    color_mode='rgb',
    class_mode = 'categorical',
    shuffle = True,
    seed = 123,
    subset = 'training', )

validation_batches = validation_datagen.flow_from_directory(
    dataset_path,
    target_size = image_size,
    batch_size = batch_size,
    color_mode='rgb',
    class_mode = 'categorical',
    shuffle = True,
    seed = 123,
    subset = 'validation', )

test_batches = validation_datagen.flow_from_directory(
    dataset_path,
    target_size = image_size,
    batch_size = batch_size,
    color_mode='rgb',
    class_mode = 'categorical',
    shuffle = True,
    seed = 123,
    subset = 'validation', )

from keras import layers
import tensorflow_hub as hub


url = "https://tfhub.dev/google/experts/bit/r50x1/in21k/bird/1"
classifier_model = url
base_model = hub.KerasLayer(classifier_model, input_shape=image_size+(3,), trainable=False)

num_of_birds = 10

model = tf.keras.Sequential([
    base_model,
    tf.keras.layers.Dense(num_of_birds, activation='softmax'),
])

model.summary()
model.compile(optimizer=keras.optimizers.Adam(), loss='categorical_crossentropy', metrics=["accuracy"],)

model.fit_generator(train_batches, epochs=5)

model.evaluate(test_batches)

import tensorflow_model_optimization as tfmot

quantize_model = tfmot.quantization.keras.quantize_model

q_aware_model = quantize_model(model)

报错信息 #

ValueError                                Traceback (most recent call last)
/usr/local/lib/python3.7/dist-packages/tensorflow_model_optimization/python/core/quantization/keras/quantize.py in quantize_apply(model, scheme)
    442   try:
--> 443     model_copy = _clone_model_with_weights(model)
    444   except ValueError as er:

12 frames
ValueError: Unknown layer: KerasLayer. Please ensure this object is passed to the `custom_objects` argument. See https://www.tensorflow.org/guide/keras/save_and_serialize#registering_the_custom_object for details.

During handling of the above exception, another exception occurred:

    ValueError                                Traceback (most recent call last)
    /usr/local/lib/python3.7/dist-packages/tensorflow_model_optimization/python/core/quantization/keras/quantize.py in quantize_apply(model, scheme)
        447         'Keras layers or objects in your model. Please specify them via '
        448         '`quantize_scope` for your calls to `quantize_model` and '
    --> 449         '`quantize_apply`. [%s].' % er)
        450 
        451   # 2. Remove QuantizeAnnotate wrappers from the layers in the model. This
    
    ValueError: Unable to clone model. This generally happens if you used custom Keras layers or objects in your model. Please specify them via `quantize_scope` for your calls to `quantize_model` and `quantize_apply`. [Unknown layer: KerasLayer. Please ensure this object is passed to the `custom_objects` argument. See https://www.tensorflow.org/guide/keras/save_and_serialize#registering_the_custom_object for details.].

解决方案 #

报错根源是TensorFlow模型优化工具tfmot无法识别TensorFlow Hub提供的KerasLayer自定义层，需要通过quantize_scope上下文管理器显式注册该层，修改量化部分的代码即可解决：

import tensorflow_model_optimization as tfmot

quantize_model = tfmot.quantization.keras.quantize_model

# 用quantize_scope注册KerasLayer，让tfmot能够识别该自定义层
with tfmot.quantization.keras.quantize_scope({'KerasLayer': hub.KerasLayer}):
    q_aware_model = quantize_model(model)

修改后，tfmot就能正确识别模型中的KerasLayer，顺利完成量化感知模型的创建。后续如果需要保存或加载这个量化模型，同样需要在quantize_scope上下文管理器中进行操作。

内容的提问来源于stack exchange，提问作者adria