如何将已有的saved_model.pb转换为EvalSavedModel？

嘿，刚好我之前折腾过这个转换，咱们把这事说清楚。先呼应你看到的TFMA指南内容：

该EvalSavedModel包含额外信息，可让TFMA以分布式方式对大量数据及用户定义分片计算模型中定义的相同评估指标。

本质上，EvalSavedModel就是在普通SavedModel基础上，多了评估专用的签名定义、指标计算图元数据，以及让TFMA能识别的分片/分布式计算逻辑。下面分两种场景给你具体实现方法：

场景1：能回溯到训练代码（最推荐） #

如果还能拿到模型训练的代码，直接在导出阶段生成EvalSavedModel是最稳妥的，不用事后折腾转换。

TensorFlow Estimator方式（兼容TF1.x/TF2.x） #

要是你用tf.estimator.Estimator训练的模型，可以直接用TFMA提供的导出工具：

import tensorflow as tf
import tensorflow_model_analysis as tfma

# 假设你已经有训练好的estimator实例
estimator = tf.estimator.DNNClassifier(...)

# 定义评估输入接收函数，要和训练的输入格式匹配
def eval_input_receiver_fn():
    serialized_tf_example = tf.compat.v1.placeholder(
        dtype=tf.string, shape=[None], name='input_example_tensor')
    # 替换成你模型实际的特征规格
    feature_spec = {'user_age': tf.io.FixedLenFeature([], tf.float32), 'label': tf.io.FixedLenFeature([], tf.int64)}
    features = tf.io.parse_example(serialized_tf_example, feature_spec)
    # 返回EvalInputReceiver，告诉TFMA输入、预测、标签的位置
    return tfma.export.EvalInputReceiver(
        features=features,
        receiver_tensors={'examples': serialized_tf_example},
        labels=features['label'])

# 导出EvalSavedModel
tfma.export.export_eval_savedmodel(
    estimator=estimator,
    export_dir_base='./eval_saved_model',
    eval_input_receiver_fn=eval_input_receiver_fn)

TF2.x Keras模型方式 #

如果是TF2.x的Keras模型，保存时额外添加评估签名就行：

import tensorflow as tf
import tensorflow_model_analysis as tfma

# 加载已训练好的Keras模型
model = tf.keras.models.load_model('./regular_saved_model')

# 定义评估专用的函数，包含指标计算
@tf.function(input_signature=[tf.TensorSpec(shape=(None, 20), dtype=tf.float32, name='features')])
def eval_step(features):
    predictions = model(features['input_features'])
    # 替换成你模型实际的指标计算逻辑
    accuracy = tf.keras.metrics.CategoricalAccuracy()(features['label'], predictions)
    return {'predictions': predictions, 'accuracy': accuracy}

# 保存为EvalSavedModel，同时保留服务签名
tf.saved_model.save(
    model,
    './eval_saved_model',
    signatures={
        'serving_default': model.signatures['serving_default'],
        'eval': eval_step.get_concrete_function()
    }
)

场景2：只有已有的saved_model.pb（无法回溯训练代码） #

这种情况得加载现有模型，手动补上评估所需的签名和指标逻辑，再重新保存。

步骤示例代码：

import tensorflow as tf
import tensorflow_model_analysis as tfma

# 加载普通SavedModel
loaded_model = tf.saved_model.load('./regular_saved_model')
# 获取默认的服务签名
infer_signature = loaded_model.signatures['serving_default']

# 提取输入输出的张量规格
input_spec = infer_signature.structured_input_signature[1]
output_pred = infer_signature.outputs['predictions']

# 定义评估函数，添加指标计算
@tf.function(input_signature=[input_spec])
def eval_fn(inputs):
    predictions = infer_signature(inputs)['predictions']
    # 假设输入里包含标签，替换成你实际的标签键
    labels = inputs['label']
    # 添加你需要的评估指标，比如损失、精准率
    loss = tf.keras.losses.SparseCategoricalCrossentropy()(labels, predictions)
    precision = tf.keras.metrics.Precision()(labels, tf.argmax(predictions, axis=1))
    return {
        'predictions': predictions,
        'cross_entropy_loss': loss,
        'precision': precision
    }

# 重新保存为EvalSavedModel
tf.saved_model.save(
    loaded_model,
    './eval_saved_model',
    signatures={
        'serving_default': infer_signature,
        'eval': eval_fn.get_concrete_function()
    }
)

验证转换是否成功 #

转换完后，用TensorFlow的命令行工具检查：

saved_model_cli show --dir ./eval_saved_model --all

如果输出里能看到eval签名，以及对应的指标输出项，就说明转换成功了，接下来就能用TFMA加载这个模型做分布式评估啦。

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