TensorFlow 2.x加载含Embedding Column的H5模型时形状不兼容报错

看起来你遇到的核心问题是DenseFeatures层的Embedding权重在加载时和当前模型结构的形状不匹配，结合你的代码和TF版本表现，我给你几个针对性的解决方案：

问题根源分析 #

你的模型使用了DenseFeatures层处理Embedding类型的特征列，这类层的权重形状完全依赖训练时生成的特征类别数量（比如NAME1W1的词汇表大小）。当你加载模型时：

1. 如果直接用load_model加载H5文件，TF无法自动复现训练时的特征列结构（比如词汇表大小、分桶数量），导致新建的模型层维度和保存的权重维度不匹配；
2. 顺序API加载失败是因为你没有先构建和训练时一致的模型结构，直接把权重加载到了空模型里；
3. TF2.2/2.3的H5保存失败是旧版本对DenseFeatures层的兼容性Bug，后续版本虽有修复，但H5本身对复杂层的支持不如原生SavedModel。

解决方案1：重建一致模型结构后加载权重 #

如果坚持用H5格式，加载前必须先1:1复现训练时的模型结构，再加载权重，而不是直接用load_model。

修改你的__loadModel函数：

def __loadModel(args):
    filepath = args.loadModel
    # 关键：用训练时完全相同的参数重建特征列
    # 注意！这里的bucketSizeGEO要和你训练保存模型时用的参数一致
    # 建议训练时把这些参数存成json文件，加载时自动读取，避免手动输入
    bucketSizeGEO = 你的训练时使用的bucketSizeGEO值
    feature_columns = preprocessing.getFutureColumns(args.data, args.zip, bucketSizeGEO, True)
    feature_layer = tf.keras.layers.DenseFeatures(feature_columns, trainable=False)
    feature_layer_inputs = preprocessing.getFeatureLayerInputs()
    feature_layer_outputs = feature_layer(feature_layer_inputs)
    output_layer = tf.keras.layers.Dense(1, activation=tf.nn.sigmoid)(feature_layer_outputs)
    # 构建和训练时完全一致的模型
    model = tf.keras.Model(inputs=[v for v in feature_layer_inputs.values()], outputs=output_layer)
    
    # 加载权重而非整个模型
    model.load_weights(filepath)
    # 必须重新编译模型
    model.compile(optimizer='sgd', loss='binary_crossentropy', metrics=['accuracy'])
    
    print("start preprocessing...")
    (_, _, test_ds) = preprocessing.getPreProcessedDatasets(args.data, args.batchSize)
    print("preprocessing completed")
    _, accuracy = model.evaluate(test_ds)
    print("Accuracy", accuracy)

解决方案2：改用SavedModel格式（推荐） #

H5格式对DenseFeatures、自定义层这类复杂结构的兼容性很差，TF原生的SavedModel格式会完整保存模型结构和所有依赖信息，是更可靠的选择。

修改训练时的保存代码： #

# 替换原来的H5保存代码
paramString = paramString + "-a{:.4f}".format(accuracy)
outputName = "logReg" + datetime.datetime.now().strftime("%Y%m%d-%H%M%S") + paramString
if args.saveModel:
    path = './saved_models/' + outputName
    model.save(path)  # 默认就是SavedModel格式，不需要指定save_format

修改加载时的代码： #

def __loadModel(args):
    filepath = args.loadModel  # 这里传入的是SavedModel的文件夹路径，不是.h5文件
    # 直接加载整个模型，无需手动重建结构
    model = tf.keras.models.load_model(filepath)
    
    print("start preprocessing...")
    (_, _, test_ds) = preprocessing.getPreProcessedDatasets(args.data, args.batchSize)
    print("preprocessing completed")
    _, accuracy = model.evaluate(test_ds)
    print("Accuracy", accuracy)

额外注意事项 #

• 参数一致性：训练和加载时的所有预处理参数（比如bucketSizeGEO、args.zip）必须完全一致，否则特征列的维度会变化，导致权重形状不匹配。建议训练时把这些参数和模型一起保存（比如用json.dump存成参数文件）。
• TF版本对齐：尽量用同一版本的TensorFlow进行训练和加载，跨版本时优先用SavedModel，它的兼容性更好。
• trainable属性一致：训练时DenseFeatures设置了trainable=False，加载时也要保持相同设置，避免权重加载异常。

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