加载模型时忽略未构建头部层？解决“预期2变量，收到0个”错误

加载通过MLflow TensorFlow flavor（Keras 3）保存的自定义ViT骨干模型时，模型包含需要的patch_embed、encoder部分，但encoder内部分层未完全构建，Keras报错：

Layer 'dense_44' expected 2 variables, but received 0 variables during loading.
Expected: ['kernel', 'bias']

需求：无需重新训练并正确保存模型，让Keras/MLflow忽略加载失败的层（无变量或未构建），不终止整个加载流程。

环境 #

• TensorFlow 2.17.0
• Keras 3.9.2
• MLflow 2.15.x
• Python 3.12

当前操作 #

已设置safe_mode=False并使用自定义对象加载模型，但仍触发错误：

import mlflow
import tensorflow as tf
from tensorflow import keras

_CUSTOM_OBJECTS = {
    "custom_object1": CustomObject1,
    # 其他自定义对象
}

model = mlflow.tensorflow.load_model(
    "runs:/<RUN_ID>/model", 
    keras_model_kwargs={"safe_mode": False, "compile": False, "custom_objects": _CUSTOM_OBJECTS},
)

尝试切换safe_mode=True/False均无效。

方法1：自定义兼容层，跳过变量检查 #

针对报错的未构建层（比如示例中的dense_44），写一个自定义子类重载构建逻辑，强制标记层已构建，避免加载时的变量校验：

class IgnoreMissingVarsDense(keras.layers.Dense):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
    
    def build(self, input_shape):
        # 直接标记层已构建，跳过变量初始化
        self.built = True

# 更新自定义对象映射，把报错的层替换成这个兼容类
_CUSTOM_OBJECTS.update({
    "dense_44": IgnoreMissingVarsDense,
    # 其他需要忽略的未构建层同理添加
})

# 重新加载模型
model = mlflow.tensorflow.load_model(
    "runs:/<RUN_ID>/model", 
    keras_model_kwargs={"safe_mode": False, "compile": False, "custom_objects": _CUSTOM_OBJECTS},
)

# 加载完成后，提取需要的patch_embed和encoder部分，丢弃无用层
patch_embed = model.get_layer("patch_embed")
encoder = model.get_layer("encoder")

# 构建只保留核心部分的新模型
inputs = keras.Input(shape=(224,224,3))
x = patch_embed(inputs)
x = encoder(x)
core_model = keras.Model(inputs=inputs, outputs=x)

方法2：手动加载配置+权重，跳过缺失项 #

绕开MLflow的直接加载逻辑，先加载模型配置，再手动加载权重并忽略缺失的变量：

import os

# 先把MLflow模型下载到本地路径
model_path = mlflow.artifacts.download_artifacts("runs:/<RUN_ID>/model")

# 加载模型配置文件
with open(os.path.join(model_path, "keras_metadata.pb"), "rb") as f:
    config = keras.saving.load_config(f)

# 把配置里的报错层替换成自定义兼容层（同方法1的IgnoreMissingVarsDense）
for layer in config["config"]["layers"]:
    if layer["class_name"] == "Dense" and layer["config"]["name"] == "dense_44":
        layer["class_name"] = "IgnoreMissingVarsDense"

# 从配置构建模型
model = keras.saving.model_from_config(config, custom_objects=_CUSTOM_OBJECTS)

# 加载权重，开启skip_missing跳过不存在的变量
model.load_weights(os.path.join(model_path, "variables/variables"), skip_missing=True)

方法3：捕获加载错误，手动修复模型结构 #

如果前两种方法仍报错，可捕获加载异常，直接提取已加载的有效层重新构建：

model_path = mlflow.artifacts.download_artifacts("runs:/<RUN_ID>/model")

try:
    model = mlflow.tensorflow.load_model(
        "runs:/<RUN_ID>/model", 
        keras_model_kwargs={"safe_mode": False, "compile": False, "custom_objects": _CUSTOM_OBJECTS},
    )
except ValueError:
    # 从配置重新构建模型，只保留需要的patch_embed和encoder
    with open(os.path.join(model_path, "keras_metadata.pb"), "rb") as f:
        config = keras.saving.load_config(f)
    # 过滤掉不需要的层
    filtered_layers = [layer for layer in config["config"]["layers"] if layer["config"]["name"] in ["patch_embed", "encoder"]]
    config["config"]["layers"] = filtered_layers
    # 构建核心模型并加载权重
    core_model = keras.saving.model_from_config(config, custom_objects=_CUSTOM_OBJECTS)
    core_model.load_weights(os.path.join(model_path, "variables/variables"), skip_missing=True)

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