R2Score Metric is incompatible with Evaluator Component

[TomsCodingCode]

R2Score Metric is incompatible with Evaluator Component

[lego0901]

Hi @TomsCodingCode, thanks for using TFX and report the issue with a concrete standalone example!

The Evaluator standard component redirects its implementation into TFMA packages. Although, I am not an expert on TFMA, at my glance, the issue arises because the tf.keras.metrics.R2Score metric stores multiple internal variables (like sum of squares, sample count, etc.) However, TFMA usually expects metrics to have a simple, single-value state for serialization and aggregation. This mismatch causes the error you encountered.

To resolve this, you can use a custom metric wrapper called R2ScoreWrapper. This wrapper encapsulates the complex internal state of R2Score and exposes only the final value to TFMA, making it compatible with TFMA's serialization and aggregation mechanisms.

class R2ScoreWrapper(tf.keras.metrics.Metric):
  def __init__(self, name="r2_score_wrapper", **kwargs):
    super().__init__(name=name, **kwargs)
    self.r2_score = tf.keras.metrics.R2Score()

  def update_state(self, y_true, y_pred, sample_weight=None):
    self.r2_score.update_state(y_true, y_pred, sample_weight)

  def result(self):
    return self.r2_score.result()

  def reset_state(self):
    self.r2_score.reset_state()

...
  linear_model.compile(
    optimizer=keras.optimizers.Adam(learning_rate=0.1),
    loss='mean_absolute_error',
    metrics=[R2ScoreWrapper()]
    # metrics=[tf.keras.metrics.R2Score()]
    )

I found the similar phenomenon appears to the metrics.F1Score too. Hope that it works as you wish. Thanks!

[TomsCodingCode]

That works perfectly fine as a workaround, thanks!

Is this a bug in tfma then or is this expected baviour?