TensorFlow 2.0 Model子类化实现ResNet模型的搭建与加载问题咨询

针对函数式网络转子类化模型时保存后加载失败的问题，核心原因是没有严格遵守tf.keras.Model子类化的序列化规范，按照以下步骤调整即可正常运行、保存、加载：

核心规范遵循要求 #

所有层的实例化操作必须放在__init__方法中，作为类的显式成员属性完成定义，严禁在call方法内临时创建层实例。这是加载失败的最主要诱因——临时创建的层不会被Keras的权重跟踪、序列化机制记录，保存的模型文件中不包含这类层的配置，加载时自然无法匹配权重结构。

具体类结构实现 #

按照你需要的网络结构，可直接参考如下标准实现：

• 在__init__方法中完成所有层的初始化
  • 加载ImageNet预训练权重的ResNet152V2骨干网络，设置include_top=False以适配自定义分类头
  • 全局平均池化层
  • 失活比例0.3的Dropout层
  • 输出维度为9的Softmax全连接分类层
• 在call方法中按数据流向依次调用已初始化的层，注意将training参数透传给所有有训练/推理模式差异的层（比如Dropout、预训练骨干层），保证训练和推理阶段的行为一致

参考完整类实现代码

import tensorflow as tf
from tensorflow.keras.applications.resnet_v2 import ResNet152V2
from tensorflow.keras.layers import GlobalAveragePooling2D, Dropout, Dense

class CustomResNet9Class(tf.keras.Model):
    def __init__(self, num_classes=9, dropout_rate=0.3, pretrained_weights="imagenet"):
        super().__init__()
        # 所有层在__init__中完成实例化，作为类属性被跟踪
        self.backbone = ResNet152V2(
            include_top=False,
            weights=pretrained_weights,
            input_shape=(224, 224, 3)
        )
        self.global_avg_pool = GlobalAveragePooling2D()
        self.dropout = Dropout(dropout_rate)
        self.output_layer = Dense(num_classes, activation="softmax")

    def call(self, inputs, training=None):
        x = self.backbone(inputs, training=training)
        x = self.global_avg_pool(x)
        x = self.dropout(x, training=training)
        return self.output_layer(x)

编译与保存加载注意事项 #

• 编译配置完全沿用你现有实现即可：使用带动量的SGD优化器、分类交叉熵损失

编译示例代码

model = CustomResNet9Class(num_classes=9, dropout_rate=0.3)
sgd_optimizer = tf.keras.optimizers.SGD(learning_rate=1e-3, momentum=0.9)
model.compile(
    optimizer=sgd_optimizer,
    loss=tf.keras.losses.CategoricalCrossentropy(),
    metrics=["accuracy"]
)

• 保存时优先使用TensorFlow SavedModel格式（直接传入文件夹路径调用model.save()即可），不要使用HDF5格式（.h5后缀），Keras子类化模型对HDF5格式的序列化支持存在较多兼容问题，极易出现加载失败
• 加载时必须将自定义模型类传入custom_objects参数，告知Keras模型对应的结构定义，否则无法解析自定义类的结构

加载示例代码

loaded_model = tf.keras.models.load_model(
    "saved_custom_resnet",
    custom_objects={"CustomResNet9Class": CustomResNet9Class}
)

常见加载失败排查点 #

• 检查__init__方法的所有入参是否设置了默认值，没有默认值的参数会导致Keras在反序列化时无法初始化类实例
• 检查call方法内是否存在临时创建层的操作，比如在call中直接写Dense(9)(x)这类没有提前在__init__中声明的层实例化逻辑
• 检查保存和加载时的网络配置是否完全一致，比如Dropout比例、输出类别数、预训练骨干的结构参数，任何配置不匹配都会导致权重形状校验失败

内容的提问来源于stack exchange，提问作者Dave-Patsy