TensorFlow中DNNLinearCombinedClassifier模型的预测获取、指标计算及模型保存问题

看起来你在使用TensorFlow Estimator的宽深模型时遇到了几个头疼的问题：自带的precision/recall指标异常、预测任务无限运行、模型保存失败，我来一步步帮你梳理解决思路~

一、先解决precision/recall为0但accuracy正常的异常 #

问题分析 #

TensorFlow Estimator默认的分类指标（precision/recall）是针对二分类场景设计的，异常大概率来自以下原因：

1. 模型默认正类别设置与你的业务不符（比如默认正类是1，但你的标签中0才是目标类别）；
2. 预测概率阈值设置过高，导致没有样本被判定为正类；
3. 标签是one-hot编码格式，但Estimator默认期望标签是整数型的类别索引；
4. 默认的eval_metric_ops计算逻辑和你的数据分布不匹配。

解决方法：自定义评估指标 #

我们可以绕过默认指标，自己定义符合业务需求的计算逻辑，在创建模型时传入：

def custom_eval_metrics(labels, predictions):
    # 获取预测的类别索引（二分类场景也可以用 probabilities[:,1] > 0.5 判定）
    pred_classes = tf.argmax(predictions['probabilities'], axis=1)
    
    # 计算自定义指标
    precision = tf.keras.metrics.Precision()(labels, pred_classes)
    recall = tf.keras.metrics.Recall()(labels, pred_classes)
    accuracy = tf.keras.metrics.Accuracy()(labels, pred_classes)
    
    return {
        'custom_precision': precision,
        'custom_recall': recall,
        'custom_accuracy': accuracy
    }

# 创建模型时指定自定义评估指标
model = tf.estimator.DNNLinearCombinedClassifier(
    linear_feature_columns=wide_columns,
    linear_optimizer=tf.keras.optimizers.legacy.Ftrl(),
    dnn_feature_columns=deep_columns,
    dnn_hidden_units=hidden_units,
    dnn_dropout=dropout_rate,
    dnn_optimizer=tf.keras.optimizers.legacy.Adagrad(),
    eval_metric_ops={'custom_metrics': custom_eval_metrics}  # 新增自定义指标
)

另外别忘了检查：你的标签数据是否是整数型的类别索引（不是one-hot编码），如果是one-hot的话，需要用tf.argmax(labels, axis=1)转成索引再计算。

二、解决predict方法无限运行的问题 #

问题根源 #

你调用model.predict()时卡住的核心原因是：你的get_input_fn没有设置num_epochs=1。TensorFlow Estimator的input_fn默认num_epochs=None（无限重复生成数据），导致predict返回的生成器会一直输出数据，Spark创建DataFrame时就会无限读取，直到你手动终止。

解决方法：修改input_fn的参数 #

确保你的输入函数设置了num_epochs=1，同时预测阶段不需要shuffle，所以也把shuffle=False加上：

# 修正你的get_input_fn（示例）
def get_input_fn(df, num_epochs=1, shuffle=False, batch_size=128):
    def input_fn():
        # 从Spark DataFrame转换为TensorFlow Dataset
        dataset = tf.data.Dataset.from_tensor_slices(dict(df))
        if shuffle:
            dataset = dataset.shuffle(buffer_size=len(df))
        # 关键：设置num_epochs=1，只生成一轮数据
        dataset = dataset.batch(batch_size).repeat(num_epochs)
        return dataset
    return input_fn

# 现在调用predict就不会无限运行了
predictions = model.predict(get_input_fn(test, num_epochs=1, shuffle=False))
predictions_df = spark.createDataFrame(predictions)

如果测试集数据量很大，建议分批处理predict结果，避免内存溢出。

三、解决模型保存失败的问题 #

你尝试的两种方法都有各自的问题，我们分别修正：

问题1：Attempt 1的ValueError（feature_columns解析冲突） #

原因

build_parsing_serving_input_receiver_fn是用来解析tf.Example格式数据的，但你的宽列中user_id是categorical_column_with_hash_bucket，对应的解析规则是VarLenFeature，而如果你的input_fn中是把user_id作为FixedLenFeature传入的，就会出现解析规则冲突。另外，你的输入来自Spark DataFrame，不是tf.Example格式，更适合用原始张量输入的serving函数。

解决方法：用build_raw_serving_input_receiver_fn

我们直接定义每个特征的占位符，匹配你的输入数据类型：

model_path = "path/to/model"

# 构建特征输入的spec：根据你的特征类型定义占位符
feature_spec = {}

# 处理宽列特征
for col in wide_columns:
    if isinstance(col, tf.feature_column.CategoricalColumn):
        # 类别型特征用int64类型，形状是[None]（支持批量输入）
        feature_spec[col.key] = tf.placeholder(dtype=tf.int64, shape=[None])

# 处理深列特征
for col in deep_columns:
    # 处理嵌入列的原始类别列
    base_col = col.categorical_column if isinstance(col, tf.feature_column.EmbeddingColumn) else col
    if isinstance(base_col, tf.feature_column.CategoricalColumn):
        feature_spec[base_col.key] = tf.placeholder(dtype=tf.int64, shape=[None])
    else:
        # 数值型特征用float32类型
        feature_spec[base_col.key] = tf.placeholder(dtype=tf.float32, shape=[None])

# 构建原始输入的serving函数
serving_input_fn = tf.estimator.export.build_raw_serving_input_receiver_fn(feature_spec)

# 导出SavedModel
export_path = model.export_saved_model(model_path, serving_input_fn)

问题2：Attempt 2的ValueError（Checkpoint不支持Estimator） #

原因

tf.train.Checkpoint和tf.saved_model.save是针对Keras模型设计的，而TensorFlow Estimator有自己的模型保存/加载机制，不能直接用这两个API。

解决方法：用Estimator原生的保存方式

Estimator在训练时会自动把checkpoint保存到你创建模型时指定的model_dir路径下（如果没指定，会生成临时目录）。你可以通过以下方式操作：

1. 加载已训练的Estimator模型：

# 从已有的model_dir加载模型
loaded_model = tf.estimator.DNNLinearCombinedClassifier(
    linear_feature_columns=wide_columns,
    dnn_feature_columns=deep_columns,
    model_dir="./your_existing_model_dir"  # 指向训练时的模型目录
)

2. 导出为SavedModel格式：
   就是上面用export_saved_model的方法，导出后可以用tf.saved_model.load加载，或者用TensorFlow Serving部署。

最后再梳理下执行顺序建议 #

1. 先修正get_input_fn的num_epochs参数，解决预测卡住的问题，这样你可以先拿到预测结果，自己用Spark或Pandas计算precision/recall，验证指标是否正常；
2. 再自定义评估指标，替换默认的eval_metric_ops，让训练时就能看到正确的指标；
3. 最后用build_raw_serving_input_receiver_fn导出SavedModel，完成模型保存。

备注：内容来源于stack exchange，提问作者Prem