TensorFlow t2.13.0-rc1加载模型遇CustomScaleLayer未知层错误求助

环境信息 #

• tensorflow版本：t2.13.0-rc1

模型训练代码 #

import os
import cv2
import numpy as np
import pandas as pd
import tensorflow as tf
import pytesseract as pt
import plotly.express as px
import matplotlib.pyplot as plt
import xml.etree.ElementTree as xet
from glob import glob
from skimage import io
from shutil import copy
from tensorflow.keras.models import Model
from tensorflow.keras.preprocessing.image import load_img, img_to_array
from sklearn.model_selection import train_test_split
from tensorflow.keras.applications import InceptionResNetV2
from tensorflow.keras.layers import Dense, Dropout, Flatten, Input
from tensorflow.keras.callbacks import TensorBoard

# 提取所有标签列的值到数组
labels = df.iloc[:,1:].values
data = []
output = []
for ind in range(len(image_path)):
    image = image_path[ind]
    img_arr = cv2.imread(image)
    h,w,d = img_arr.shape
    # 预处理
    load_image = load_img(image,target_size=(224,224))
    load_image_arr = img_to_array(load_image)
    norm_load_image_arr = load_image_arr/255.0 # 归一化
    # 标签归一化
    xmin,xmax,ymin,ymax = labels[ind]
    nxmin,nxmax = xmin/w,xmax/w
    nymin,nymax = ymin/h,ymax/h
    label_norm = (nxmin,nxmax,nymin,nymax) # 归一化后的标签
    # 添加到列表
    data.append(norm_load_image_arr)
    output.append(label_norm)

# 转换为数组
X = np.array(data,dtype=np.float32)
y = np.array(output,dtype=np.float32)

# 用sklearn划分训练集和测试集
x_train, x_test, y_train, y_test = train_test_split(X, y, train_size=0.8, random_state=0)
x_train.shape,x_test.shape,y_train.shape,y_test.shape

# 构建模型
inception_resnet = InceptionResNetV2(weights="imagenet",include_top=False, input_tensor=Input(shape=(224,224,3)))
# ---------------------
headmodel = inception_resnet.output
headmodel = Flatten()(headmodel)
headmodel = Dense(500,activation="relu")(headmodel)
headmodel = Dense(250,activation="relu")(headmodel)
headmodel = Dense(4,activation='sigmoid')(headmodel)


# ---------- 定义完整模型
model = Model(inputs=inception_resnet.input,outputs=headmodel)

model.compile(loss='mse', optimizer=tf.keras.optimizers.legacy.Adam(learning_rate=1e-4))
model.summary()

模型摘要片段 #

# 模型摘要截取（仅展示最后几行）

 conv2d_605 (Conv2D)         (None, 5, 5, 192)            399360    ['block8_9_ac[0][0]']         
                                                                                                  
 conv2d_608 (Conv2D)         (None, 5, 5, 256)            172032    ['activation_607[0][0]']      
                                                                                                  
 batch_normalization_605 (BatchNormalization)            (None, 5, 5, 192)            576       ['conv2d_605[0][0]']          
                                                                                                  
 batch_normalization_608 (BatchNormalization)            (None, 5, 5, 256)            768       ['conv2d_608[0][0]']          
                                                                                                  
 activation_605 (Activation)            (None, 5, 5, 192)            0         ['batch_normalization_605[0][0]']                            
                                                                                                  
 activation_608 (Activation)            (None, 5, 5, 256)            0         ['batch_normalization_608[0][0]']                            
                                                                                                  
 block8_10_mixed (Concatenate)            (None, 5, 5, 448)            0         ['activation_605[0][0]',      
                                                                 'activation_608[0][0]']      
                                                                                                  
 block8_10_conv (Conv2D)     (None, 5, 5, 2080)           933920    ['block8_10_mixed[0][0]']     
                                                                                                  
 custom_scale_layer_119 (CustomScaleLayer)            (None, 5, 5, 2080)           0         ['block8_9_ac[0][0]',         
                                                                 'block8_10_conv[0][0]']      
                                                                                                  
 conv_7b (Conv2D)            (None, 5, 5, 1536)           3194880   ['custom_scale_layer_119[0][0]']                             
                                                                                                  
 conv_7b_bn (BatchNormalization)            (None, 5, 5, 1536)           4608      ['conv_7b[0][0]']              
                                                                                                  
 conv_7b_ac (Activation)     (None, 5, 5, 1536)           0         ['conv_7b_bn[0][0]']          
                                                                                                  
 flatten_2 (Flatten)         (None, 38400)                0         ['conv_7b_ac[0][0]']          
                                                                                                  
 dense_6 (Dense)             (None, 500)                  19200500   ['flatten_2[0][0]']           
                                                                                                  
 dense_7 (Dense)             (None, 250)                  125250    ['dense_6[0][0]']              
                                                                                                  
 dense_8 (Dense)             (None, 4)                    1004      ['dense_7[0][0]']              
                                                                                                  
==================================================================================================
Total params: 73663490 (281.00 MB)
Trainable params: 73602946 (280.77 MB)
Non-trainable params: 60544 (236.50 KB)

模型保存代码 #

tfb = TensorBoard('object_detection2')
history = model.fit(x=x_train,y=y_train,batch_size=10,epochs=180,
                    validation_data=(x_test,y_test),callbacks=[tfb])

model.save('./object_detection.h5')

模型加载代码及报错 #

模型训练完成后，执行以下加载代码：

model = tf.keras.models.load_model('./object_detection.h5')
print('Model loaded Sucessfully')

触发如下错误：

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[30], line 2
      1 # Load model
----> 2 model = tf.keras.models.load_model('./object_detection.h5')
      3 print('Model loaded Sucessfully')

File ~/anaconda3/lib/python3.10/site-packages/keras/src/saving/saving_api.py:238, in load_model(filepath, custom_objects, compile, safe_mode, **kwargs)
    230     return saving_lib.load_model(
    231         filepath,
    232         custom_objects=custom_objects,
    233         compile=compile,
    234         safe_mode=safe_mode,
    235     )
    237 # Legacy case.
--> 238 return legacy_sm_saving_lib.load_model(
    239     filepath, custom_objects=custom_objects, compile=compile, **kwargs
    240 )

File ~/anaconda3/lib/python3.10/site-packages/keras/src/utils/traceback_utils.py:70, in filter_traceback.<locals>.error_handler(*args, **kwargs)
     67     filtered_tb = _process_traceback_frames(e.__traceback__)
     68     # To get the full stack trace, call:
     69     # `tf.debugging.disable_traceback_filtering()`
--> 70     raise e.with_traceback(filtered_tb) from None
     71 finally:
     72     del filtered_tb

File ~/anaconda3/lib/python3.10/site-packages/keras/src/saving/legacy/serialization.py:365, in class_and_config_for_serialized_keras_object(config, module_objects, custom_objects, printable_module_name)
    361 cls = object_registration.get_registered_object(
    362     class_name, custom_objects, module_objects
    363 )
    364 if cls is None:
--> 365     raise ValueError(
    366         f"Unknown {printable_module_name}: '{class_name}'. "
    367         "Please ensure you are using a `keras.utils.custom_object_scope` "
    368         "and that this object is included in the scope. See "
    369         "https://www.tensorflow.org/guide/keras/save_and_serialize"
    370         "#registering_the_custom_object for details."
    371     )
    373 cls_config = config["config"]
    374 # Check if `cls_config` is a list. If it is a list, return the class and the
    375 # associated class configs for recursively deserialization. This case will
    376 # happen on the old version of sequential model (e.g. `keras_version` ==
    377 # "2.0.6"), which is serialized in a different structure, for example
    378 # "{'class_name': 'Sequential',
    379 #   'config': [{'class_name': 'Embedding', 'config': ...}, {}, ...]}".

ValueError: Unknown layer: 'CustomScaleLayer'. Please ensure you are using a `keras.utils.custom_object_scope` and that this object is included in the scope. See https://www.tensorflow.org/guide/keras/save_and_serialize#registering_the_custom_object for details.

问题 #

如何解决该模型加载错误，成功加载训练好的模型？

解决方法 #

错误原因是CustomScaleLayer是InceptionResNetV2内部的自定义层，Keras加载.h5格式模型时无法自动识别该层，以下是两种可行的解决方式：

方式一：加载时传入自定义层 #

从TensorFlow的InceptionResNetV2模块中导入CustomScaleLayer，并在加载模型时指定为自定义对象：

from tensorflow.keras.applications.inception_resnet_v2 import CustomScaleLayer
from tensorflow.keras.models import load_model

model = load_model('./object_detection.h5', custom_objects={'CustomScaleLayer': CustomScaleLayer})
print('Model loaded Successfully')

方式二：改用SavedModel格式保存（推荐） #

SavedModel是TensorFlow官方推荐的模型保存格式，能更好地兼容复杂模型结构，避免自定义层识别问题。修改保存代码如下：

# 训练完成后保存为SavedModel格式
model.save('./object_detection_model')

# 加载时无需指定自定义对象
model = tf.keras.models.load_model('./object_detection_model')
print('Model loaded Successfully')

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