如何创建一个在 tf.estimator.train_and_evaluate 评估步骤中持续存在的变量？

Question

TLDR：如何创建一个变量来保存用于计算自定义指标的混淆矩阵，并在所有评估步骤中累积值？

我have implemented 自定义指标在tf.estimator.train_and_evaluation 管道中使用，混淆矩阵是所有指标的关键。我的目标是让这个混淆矩阵在多个评估步骤中持续存在，以便跟踪学习进度。

在变量范围内使用get_variable 不起作用，因为它不会将变量保存到检查点（或者看起来如此）。

这不起作用：

    @property
    def confusion_matrix(self):
        with tf.variable_scope(
            f"{self.name}-{self.metric_type}", reuse=tf.AUTO_REUSE
        ):
            confusion_matrix = tf.get_variable(
                name="confusion-matrix",
                initializer=tf.zeros(
                    [self.n_classes, self.n_classes],
                    dtype=tf.float32,
                    name=f"{self.name}/{self.metric_type}-confusion-matrix",
                ),
                aggregation=tf.VariableAggregation.SUM,
            )
        return confusion_matrix

只需将矩阵保存为类属性即可，但显然不会在多个步骤中持续存在：

        self.confusion_matrix = tf.zeros(
            [self.n_classes, self.n_classes],
            dtype=tf.float32,
            name=f"{self.name}/{self.metric_type}-confusion-matrix",
        )

您可以查看完整示例here。

我希望这个混淆矩阵在评估期间从头到尾持续存在，但我不需要在最终的 SavedModel 中使用它。您能告诉我如何实现这一目标吗？我需要将矩阵保存到外部文件，还是有更好的方法？

Answer 1

您可以定义自定义指标：

def confusion_matrix(labels, predictions):
    matrix = ... # confusion matrix calculation
    mean, update_op = tf.metrics.mean_tensor(matrix)
    # do something with mean if needed
    return {'confusion_matrix': (mean, udpate_op)}

然后将其添加到您的estimator：

estimator = tf.estimator.add_metrics(estimator, confusion_matrix)

如果您需要 sum 而不是 meen，您可以从 tf.metrics.mean_tensor 实现中获得洞察