diff --git a/src/metrax/__init__.py b/src/metrax/__init__.py index 5b7b151..d233450 100644 --- a/src/metrax/__init__.py +++ b/src/metrax/__init__.py @@ -12,6 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +"""Metrax metrics collection.""" + from metrax import audio_metrics from metrax import base from metrax import classification_metrics @@ -26,6 +28,8 @@ Average = base.Average AveragePrecisionAtK = ranking_metrics.AveragePrecisionAtK BLEU = nlp_metrics.BLEU +BinaryAccuracy = classification_metrics.BinaryAccuracy +CategoricalAccuracy = classification_metrics.CategoricalAccuracy CosineSimilarity = image_metrics.CosineSimilarity DCGAtK = ranking_metrics.DCGAtK Dice = image_metrics.Dice @@ -48,6 +52,7 @@ RougeL = nlp_metrics.RougeL RougeN = nlp_metrics.RougeN SNR = audio_metrics.SNR +SparseCategoricalAccuracy = classification_metrics.SparseCategoricalAccuracy SpearmanRankCorrelation = regression_metrics.SpearmanRankCorrelation SSIM = image_metrics.SSIM WER = nlp_metrics.WER @@ -60,6 +65,8 @@ "Average", "AveragePrecisionAtK", "BLEU", + "BinaryAccuracy", + "CategoricalAccuracy", "CosineSimilarity", "DCGAtK", "Dice", @@ -77,6 +84,7 @@ "RMSE", "RMSLE", "RSQUARED", + "SparseCategoricalAccuracy", "SpearmanRankCorrelation", "Recall", "RecallAtK", diff --git a/src/metrax/classification_metrics.py b/src/metrax/classification_metrics.py index cf3a070..12872cf 100644 --- a/src/metrax/classification_metrics.py +++ b/src/metrax/classification_metrics.py @@ -110,6 +110,125 @@ def from_model_output( ) +@flax.struct.dataclass +class BinaryAccuracy(base.Average): + r"""Computes binary classification accuracy for predictions and labels. + + This metric calculates the proportion of correct predictions by comparing + `predictions >= threshold` and `labels` element-wise. It is the ratio of the + sum of weighted correct predictions to the sum of all corresponding weights. + If no `sample_weights` are provided, weights default to 1 for each element. + """ + + @classmethod + def from_model_output( + cls, + predictions: jax.Array, + labels: jax.Array, + sample_weights: jax.Array | None = None, + threshold: float = 0.5, + ) -> 'BinaryAccuracy': + """Updates the metric state with new `predictions` and `labels`. + + Args: + predictions: JAX array of predicted values. + labels: JAX array of true binary values (0 or 1). + sample_weights: Optional JAX array of sample weights. + threshold: The threshold parameter used to convert predicted probabilities + to binary decisions. + + Returns: + An updated instance of `BinaryAccuracy` metric. + """ + correct = (predictions >= threshold) == labels + count = jnp.ones_like(labels, dtype=jnp.int32) + if sample_weights is not None: + correct = correct * sample_weights + count = count * sample_weights + return cls( + total=correct.sum(), + count=count.sum(), + ) + + +@flax.struct.dataclass +class CategoricalAccuracy(base.Average): + r"""Computes accuracy for one-hot categorical classification models. + + This metric calculates the frequency with which the predicted class matches + the true class by comparing `argmax(predictions, axis=-1)` and + `argmax(labels, axis=-1)`. If no `sample_weights` are provided, weights + default to 1 for each element. + """ + + @classmethod + def from_model_output( + cls, + predictions: jax.Array, + labels: jax.Array, + sample_weights: jax.Array | None = None, + ) -> 'CategoricalAccuracy': + """Updates the metric state with new `predictions` and `labels`. + + Args: + predictions: JAX array of predicted probability distributions or logits. + labels: JAX array of one-hot encoded true class labels. + sample_weights: Optional JAX array of sample weights. + + Returns: + An updated instance of `CategoricalAccuracy` metric. + """ + correct = jnp.argmax(predictions, axis=-1) == jnp.argmax(labels, axis=-1) + count = jnp.ones_like(correct, dtype=jnp.int32) + if sample_weights is not None: + correct = correct * sample_weights + count = count * sample_weights + return cls( + total=correct.sum(), + count=count.sum(), + ) + + +@flax.struct.dataclass +class SparseCategoricalAccuracy(base.Average): + r"""Computes accuracy for sparse categorical classification models. + + This metric calculates the frequency with which the predicted class matches + the integer target class by comparing `argmax(predictions, axis=-1)` and + `labels`. If no `sample_weights` are provided, weights default to 1 for + each element. + """ + + @classmethod + def from_model_output( + cls, + predictions: jax.Array, + labels: jax.Array, + sample_weights: jax.Array | None = None, + ) -> 'SparseCategoricalAccuracy': + """Updates the metric state with new `predictions` and `labels`. + + Args: + predictions: JAX array of predicted probability distributions or logits. + labels: JAX array of ground truth integer class labels. + sample_weights: Optional JAX array of sample weights. + + Returns: + An updated instance of `SparseCategoricalAccuracy` metric. + """ + if labels.ndim == predictions.ndim: + labels = jnp.squeeze(labels, axis=-1) + correct = jnp.argmax(predictions, axis=-1) == labels + count = jnp.ones_like(labels, dtype=jnp.int32) + if sample_weights is not None: + correct = correct * sample_weights + count = count * sample_weights + return cls( + total=correct.sum(), + count=count.sum(), + ) + + @flax.struct.dataclass class Precision(clu_metrics.Metric): r"""Computes precision for binary classification given `predictions` and `labels`. diff --git a/src/metrax/classification_metrics_test.py b/src/metrax/classification_metrics_test.py index c9255c3..9692450 100644 --- a/src/metrax/classification_metrics_test.py +++ b/src/metrax/classification_metrics_test.py @@ -47,6 +47,16 @@ [0.5, 1, 0, 0, 0, 0, 0, 0], (BATCHES, 1), ).astype(np.float32) +NUM_CLASSES = 5 +MC_LABELS_INT = np.random.randint( + 0, NUM_CLASSES, size=(BATCHES, BATCH_SIZE) +).astype(np.int32) +MC_LABELS_OH = np.eye(NUM_CLASSES)[MC_LABELS_INT].astype(np.float32) +MC_PREDS = np.random.uniform(size=(BATCHES, BATCH_SIZE, NUM_CLASSES)).astype( + np.float32 +) +MC_PREDS_F16 = MC_PREDS.astype(jnp.float16) +MC_PREDS_BF16 = MC_PREDS.astype(jnp.bfloat16) class ClassificationMetricsTest(parameterized.TestCase): @@ -89,6 +99,24 @@ def test_fbeta_empty(self): self.assertEqual(m.false_negatives, jnp.array(0, jnp.float32)) self.assertEqual(m.beta, 1.0) + def test_binary_accuracy_empty(self): + """Tests the `empty` method of `BinaryAccuracy`.""" + m = metrax.BinaryAccuracy.empty() + self.assertEqual(m.total, jnp.array(0, jnp.float32)) + self.assertEqual(m.count, jnp.array(0, jnp.int32)) + + def test_categorical_accuracy_empty(self): + """Tests the `empty` method of `CategoricalAccuracy`.""" + m = metrax.CategoricalAccuracy.empty() + self.assertEqual(m.total, jnp.array(0, jnp.float32)) + self.assertEqual(m.count, jnp.array(0, jnp.int32)) + + def test_sparse_categorical_accuracy_empty(self): + """Tests the `empty` method of `SparseCategoricalAccuracy`.""" + m = metrax.SparseCategoricalAccuracy.empty() + self.assertEqual(m.total, jnp.array(0, jnp.float32)) + self.assertEqual(m.count, jnp.array(0, jnp.int32)) + @parameterized.named_parameters( ('basic_f16', OUTPUT_LABELS, OUTPUT_PREDS_F16, SAMPLE_WEIGHTS), ('basic_f32', OUTPUT_LABELS, OUTPUT_PREDS_F32, SAMPLE_WEIGHTS), @@ -120,6 +148,99 @@ def test_accuracy(self, y_true, y_pred, sample_weights): atol=atol, ) + @parameterized.named_parameters( + ('f16', OUTPUT_LABELS, OUTPUT_PREDS_F16, SAMPLE_WEIGHTS, 0.5), + ('high_f16', OUTPUT_LABELS, OUTPUT_PREDS_F16, SAMPLE_WEIGHTS, 0.7), + ('low_f16', OUTPUT_LABELS, OUTPUT_PREDS_F16, SAMPLE_WEIGHTS, 0.1), + ('f32', OUTPUT_LABELS, OUTPUT_PREDS_F32, SAMPLE_WEIGHTS, 0.5), + ('high_f32', OUTPUT_LABELS, OUTPUT_PREDS_F32, SAMPLE_WEIGHTS, 0.7), + ('low_f32', OUTPUT_LABELS, OUTPUT_PREDS_F32, SAMPLE_WEIGHTS, 0.1), + ('bf16', OUTPUT_LABELS, OUTPUT_PREDS_BF16, SAMPLE_WEIGHTS, 0.5), + ('bs_one', OUTPUT_LABELS_BS1, OUTPUT_PREDS_BS1, None, 0.5), + ) + def test_binary_accuracy(self, y_true, y_pred, sample_weights, threshold): + """Test that `BinaryAccuracy` metric computes correct values.""" + if sample_weights is None: + sample_weights = np.ones_like(y_true) + metrax_metric = metrax.BinaryAccuracy.empty() + keras_metric = keras.metrics.BinaryAccuracy(threshold=threshold) + for labels, logits, weights in zip(y_true, y_pred, sample_weights): + update = metrax.BinaryAccuracy.from_model_output( + predictions=logits, + labels=labels, + sample_weights=weights, + threshold=threshold, + ) + metrax_metric = metrax_metric.merge(update) + keras_metric.update_state(labels, logits, weights) + + rtol = 1e-2 if y_pred.dtype in (jnp.float16, jnp.bfloat16) else 1e-5 + atol = 1e-2 if y_pred.dtype in (jnp.float16, jnp.bfloat16) else 1e-5 + np.testing.assert_allclose( + metrax_metric.compute(), + keras_metric.result(), + rtol=rtol, + atol=atol, + ) + + @parameterized.named_parameters( + ('f16', MC_LABELS_OH, MC_PREDS_F16, SAMPLE_WEIGHTS), + ('f32', MC_LABELS_OH, MC_PREDS, SAMPLE_WEIGHTS), + ('bf16', MC_LABELS_OH, MC_PREDS_BF16, None), + ) + def test_categorical_accuracy(self, y_true, y_pred, sample_weights): + """Test that `CategoricalAccuracy` metric computes correct values.""" + if sample_weights is None: + sample_weights = np.ones((BATCHES, BATCH_SIZE), dtype=np.float32) + metrax_metric = metrax.CategoricalAccuracy.empty() + keras_metric = keras.metrics.CategoricalAccuracy() + for labels, logits, weights in zip(y_true, y_pred, sample_weights): + update = metrax.CategoricalAccuracy.from_model_output( + predictions=logits, + labels=labels, + sample_weights=weights, + ) + metrax_metric = metrax_metric.merge(update) + keras_metric.update_state(labels, logits, weights) + + rtol = 1e-2 if y_pred.dtype in (jnp.float16, jnp.bfloat16) else 1e-5 + atol = 1e-2 if y_pred.dtype in (jnp.float16, jnp.bfloat16) else 1e-5 + np.testing.assert_allclose( + metrax_metric.compute(), + keras_metric.result(), + rtol=rtol, + atol=atol, + ) + + @parameterized.named_parameters( + ('f16', MC_LABELS_INT, MC_PREDS_F16, SAMPLE_WEIGHTS), + ('f32', MC_LABELS_INT, MC_PREDS, SAMPLE_WEIGHTS), + ('bf16', MC_LABELS_INT, MC_PREDS_BF16, None), + ) + def test_sparse_categorical_accuracy(self, y_true, y_pred, sample_weights): + """Test that `SparseCategoricalAccuracy` computes correct values.""" + if sample_weights is None: + sample_weights = np.ones((BATCHES, BATCH_SIZE), dtype=np.float32) + metrax_metric = metrax.SparseCategoricalAccuracy.empty() + keras_metric = keras.metrics.SparseCategoricalAccuracy() + for labels, logits, weights in zip(y_true, y_pred, sample_weights): + update = metrax.SparseCategoricalAccuracy.from_model_output( + predictions=logits, + labels=labels, + sample_weights=weights, + ) + metrax_metric = metrax_metric.merge(update) + keras_metric.update_state(labels, logits, weights) + + rtol = 1e-2 if y_pred.dtype in (jnp.float16, jnp.bfloat16) else 1e-5 + atol = 1e-2 if y_pred.dtype in (jnp.float16, jnp.bfloat16) else 1e-5 + np.testing.assert_allclose( + metrax_metric.compute(), + keras_metric.result(), + rtol=rtol, + atol=atol, + ) + @parameterized.named_parameters( ('basic_f16', OUTPUT_LABELS, OUTPUT_PREDS_F16, 0.5), ('high_threshold_f16', OUTPUT_LABELS, OUTPUT_PREDS_F16, 0.7), diff --git a/src/metrax/nnx/__init__.py b/src/metrax/nnx/__init__.py index df30ccf..77a7da4 100644 --- a/src/metrax/nnx/__init__.py +++ b/src/metrax/nnx/__init__.py @@ -12,6 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +"""Metrax NNX metrics collection.""" + from metrax.nnx import nnx_metrics AUCPR = nnx_metrics.AUCPR @@ -20,6 +22,8 @@ Average = nnx_metrics.Average AveragePrecisionAtK = nnx_metrics.AveragePrecisionAtK BLEU = nnx_metrics.BLEU +BinaryAccuracy = nnx_metrics.BinaryAccuracy +CategoricalAccuracy = nnx_metrics.CategoricalAccuracy CosineSimilarity = nnx_metrics.CosineSimilarity DCGAtK = nnx_metrics.DCGAtK Dice = nnx_metrics.Dice @@ -42,6 +46,7 @@ RougeL = nnx_metrics.RougeL RougeN = nnx_metrics.RougeN SNR = nnx_metrics.SNR +SparseCategoricalAccuracy = nnx_metrics.SparseCategoricalAccuracy SpearmanRankCorrelation = nnx_metrics.SpearmanRankCorrelation SSIM = nnx_metrics.SSIM WER = nnx_metrics.WER @@ -53,6 +58,8 @@ "Average", "AveragePrecisionAtK", "BLEU", + "BinaryAccuracy", + "CategoricalAccuracy", "CosineSimilarity", "DCGAtK", "Dice", @@ -73,6 +80,7 @@ "RougeL", "RougeN", "SNR", + "SparseCategoricalAccuracy", "SpearmanRankCorrelation", "SSIM", "WER", diff --git a/src/metrax/nnx/nnx_metrics.py b/src/metrax/nnx/nnx_metrics.py index 5177d46..82524c2 100644 --- a/src/metrax/nnx/nnx_metrics.py +++ b/src/metrax/nnx/nnx_metrics.py @@ -12,6 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +"""NNX wrapper metric implementations for Metrax.""" + import metrax from metrax.nnx import nnx_wrapper @@ -80,12 +82,14 @@ class Dice(NnxWrapper): def __init__(self): super().__init__(metrax.Dice) + class FBetaScore(NnxWrapper): """An NNX class for the Metrax metric FBetaScore.""" def __init__(self): super().__init__(metrax.FBetaScore) + class IoU(NnxWrapper): """An NNX class for the Metrax metric IoU.""" @@ -231,3 +235,24 @@ class WER(NnxWrapper): def __init__(self): super().__init__(metrax.WER) + + +class BinaryAccuracy(NnxWrapper): + """An NNX class for the Metrax metric BinaryAccuracy.""" + + def __init__(self): + super().__init__(metrax.BinaryAccuracy) + + +class CategoricalAccuracy(NnxWrapper): + """An NNX class for the Metrax metric CategoricalAccuracy.""" + + def __init__(self): + super().__init__(metrax.CategoricalAccuracy) + + +class SparseCategoricalAccuracy(NnxWrapper): + """An NNX class for the Metrax metric SparseCategoricalAccuracy.""" + + def __init__(self): + super().__init__(metrax.SparseCategoricalAccuracy)