TFX流水线Evaluator组件无法获取Fare标签，如何配置以计算R2得分？

核心问题分析 #

当前模型仅导出了服务用的serving_default签名，且预测结果使用通用的predictions键，Evaluator无法将其与数据中的fare标签对应；同时缺少专门用于评估的签名，导致Evaluator无法正确关联标签与预测值。

修改步骤 #

1. 调整预测输出键与标签列名一致 #

将serving_step的返回键从predictions改为fare，让预测结果的键直接匹配标签列名：

@tf.function(input_signature=[
    tf.TensorSpec(shape=(None, 1), dtype=tf.float32, name="trip_miles"),
    tf.TensorSpec(shape=(None, 1), dtype=tf.int64, name="trip_seconds"),
    tf.TensorSpec(shape=(None, 1), dtype=tf.int64, name="trip_start_hour"),
    tf.TensorSpec(shape=(None, 1), dtype=tf.int64, name="company_size_encoded"),
    tf.TensorSpec(shape=(None, 1), dtype=tf.int64, name="pickup_community_area"),
    tf.TensorSpec(shape=(None, 1), dtype=tf.int64, name="dropoff_community_area"),
])
def serving_step(trip_miles, trip_seconds, trip_start_hour, company_size_encoded,
                 pickup_community_area, dropoff_community_area):
    inputs = [
        trip_miles, trip_seconds, trip_start_hour, 
        company_size_encoded, pickup_community_area, dropoff_community_area
    ]
    predictions = model(inputs, training=False)
    # 将返回键改为与标签列名一致的"fare"
    return {"fare": predictions}

2. 添加评估专用签名 #

新增eval_step函数，创建评估签名，让Evaluator能明确获取标签与预测值：

@tf.function(input_signature=[
    tf.TensorSpec(shape=(None, 1), dtype=tf.float32, name="trip_miles"),
    tf.TensorSpec(shape=(None, 1), dtype=tf.int64, name="trip_seconds"),
    tf.TensorSpec(shape=(None, 1), dtype=tf.int64, name="trip_start_hour"),
    tf.TensorSpec(shape=(None, 1), dtype=tf.int64, name="company_size_encoded"),
    tf.TensorSpec(shape=(None, 1), dtype=tf.int64, name="pickup_community_area"),
    tf.TensorSpec(shape=(None, 1), dtype=tf.int64, name="dropoff_community_area"),
    # 加入标签列的TensorSpec
    tf.TensorSpec(shape=(None, 1), dtype=tf.float32, name="fare"),
])
def eval_step(trip_miles, trip_seconds, trip_start_hour, company_size_encoded,
              pickup_community_area, dropoff_community_area, fare):
    inputs = [
        trip_miles, trip_seconds, trip_start_hour, 
        company_size_encoded, pickup_community_area, dropoff_community_area
    ]
    predictions = model(inputs, training=False)
    # 返回预测值和真实标签，方便Evaluator计算指标
    return {"fare": predictions, "labels": fare}

3. 更新签名字典并保存模型 #

将评估签名加入到signatures中，确保模型导出时包含服务和评估两种签名：

signatures = {
    "serving_default": serving_step.get_concrete_function(),
    "eval_default": eval_step.get_concrete_function(),  # 添加评估签名
}

# 保存模型
tf.keras.models.save_model(
    model, 
    fn_args.serving_model_dir,
    signatures=signatures
)

4. 配置Evaluator组件指定标签 #

在定义Evaluator时，明确指定标签键为fare，并配置R2得分的评估指标：

from tfx.components import Evaluator
from tfx.proto import evaluator_pb2

evaluator = Evaluator(
    examples=example_gen.outputs['examples'],
    model=trainer.outputs['model'],
    eval_config=evaluator_pb2.EvalConfig(
        model_specs=[evaluator_pb2.ModelSpec(label_key='fare')],
        metrics_specs=[
            evaluator_pb2.MetricsSpec(
                metrics=[
                    evaluator_pb2.MetricConfig(class_name='R2Score',
                                              config={'name': 'r2_score'})
                ]
            )
        ]
    )
)

原理说明 #

• 调整预测输出键为fare，让Evaluator能直接将预测结果与数据中的标签列匹配；
• 添加评估签名后，Evaluator可以通过该签名明确获取真实标签和预测值，避免标签匹配错误；
• 在Evaluator配置中指定label_key='fare'，进一步明确评估时使用的标签列。

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