如何在TensorFlow Estimator的eval_metric_ops中实现中位数并在TensorBoard展示

我来帮你搞定这个问题——你现在的方法能算出中位数，但没法在TensorBoard里实时查看，核心是要实现一个符合Estimator eval_metric_ops要求的流式指标，也就是返回(结果张量, 更新操作)的元组，这样TensorFlow就能自动把指标写入TensorBoard了。

方法一：手动实现自定义流式中位数 #

TensorFlow原生确实没内置这个功能，但我们可以用变量缓存所有评估样本，动态计算中位数，代码如下：

1. 编写自定义指标函数 #

import tensorflow as tf

def streaming_median(values, name=None):
    """返回符合Estimator要求的(中位数张量, 更新操作)元组"""
    with tf.compat.v1.variable_scope(name, "streaming_median", [values]):
        # 定义变量缓存所有评估过程中的样本，初始为空
        all_samples = tf.compat.v1.get_variable(
            "all_samples",
            shape=[0],
            dtype=values.dtype,
            initializer=tf.compat.v1.zeros_initializer(),
            trainable=False  # 这个变量不需要参与训练
        )
        
        # 更新操作：把当前batch的样本追加到缓存里
        updated_samples = tf.concat([all_samples, tf.reshape(values, [-1])], axis=0)
        update_op = tf.compat.v1.assign(all_samples, updated_samples)
        
        # 计算中位数：用percentile取50分位数，插值选midpoint保证结果准确
        median_tensor = tf.contrib.distributions.percentile(all_samples, 50.0, interpolation='midpoint')
        
        return median_tensor, update_op

2. 在model_function中替换指标 #

把原来的tf.metrics.mean换成这个自定义函数就行：

if mode == tf.estimator.ModeKeys.EVAL:
    return tf.estimator.EstimatorSpec(
        mode=mode,
        loss=loss,
        eval_metric_ops={"accuracy_median": streaming_median(my_tensor)}
    )

3. 评估时的小提示 #

评估的时候记得设置steps=None，这样能覆盖整个评估数据集，保证中位数计算的准确性。而且Estimator每次评估前会自动重置指标变量，不用你手动清理缓存。

方法二：用TensorFlow Addons简化实现 #

如果你能安装tensorflow-addons库，里面有现成的Median指标类，用法更省心：

import tensorflow_addons as tfa

if mode == tf.estimator.ModeKeys.EVAL:
    median_metric = tfa.metrics.Median()
    median_metric.update_state(my_tensor)
    return tf.estimator.EstimatorSpec(
        mode=mode,
        loss=loss,
        eval_metric_ops={"accuracy_median": median_metric.result()}
    )

这个类底层也是通过缓存样本计算中位数，和我们手动实现的逻辑一致，只是封装得更友好。

为什么这样能在TensorBoard看到指标？ #

因为我们返回的中位数结果是TensorFlow计算图里的张量，更新操作会在评估过程中逐步执行，Estimator会自动把这个指标的结果写入TensorBoard日志，这样你就能在训练/评估过程中实时看到中位数的变化了，完美解决你之前的问题～

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