More Metrics

examples/cnn.py

@@ -16,7 +16,6 @@
 
 Batch = tp.Mapping[str, np.ndarray]
 Model = tx.Sequential
-Metric = tx.metrics.Accuracy
 np.random.seed(420)
 
 
@@ -31,7 +30,7 @@ def init_step(
 
 
 @jax.jit
-def reset_step(metric: Metric) -> Metric:
+def reset_step(metric: tx.Metric) -> tx.Metric:
     metric.reset()
     return metric
 
@@ -59,10 +58,10 @@ def loss_fn(
 def train_step(
     model: Model,
     optimizer: tx.Optimizer,
-    metric: Metric,
+    metric: tx.Metric,
     x: jnp.ndarray,
     y: jnp.ndarray,
-) -> tp.Tuple[jnp.ndarray, Model, tx.Optimizer, Metric]:
+) -> tp.Tuple[jnp.ndarray, Model, tx.Optimizer, tx.Metric]:
     print("JITTTTING")
     params = model.parameters()
 
@@ -79,8 +78,8 @@ def train_step(
 
 @jax.jit
 def test_step(
-    model: Model, metric: Metric, x: jnp.ndarray, y: jnp.ndarray
-) -> tp.Tuple[jnp.ndarray, Metric]:
+    model: Model, metric: tx.Metric, x: jnp.ndarray, y: jnp.ndarray
+) -> tp.Tuple[jnp.ndarray, tx.Metric]:
 
     loss, (model, y_pred) = loss_fn(model, model, x, y)
 

tests/metrics/test_accuracy.py

@@ -5,8 +5,8 @@
 import pytest
 
 import treex as tx
-from treex.metrics.classification.accuracy import Accuracy
-from treex.metrics.utilities.enums import DataType
+from treex.metrics.accuracy import Accuracy
+from treex.metrics.utils import DataType
 
 
 class TestAccuracy:
@@ -21,15 +21,15 @@ def f(m, y_true, y_pred):
             return m
 
         metric = Accuracy(num_classes=10)
-        y_true = jnp.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
-        y_pred = jnp.array([0, 1, 2, 3, 0, 5, 6, 7, 0, 9])
+        y_true = jnp.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])[None, None, None, :]
+        y_pred = jnp.array([0, 1, 2, 3, 0, 5, 6, 7, 0, 9])[None, None, None, :]
 
         metric = f(metric, y_true, y_pred)
         assert N == 1
         assert metric.compute() == 0.8
 
         metric = f(metric, y_true, y_pred)
-        # assert N == 1
+        assert N == 1
         assert metric.compute() == 0.8
 
     def test_logits_preds(self):
@@ -54,7 +54,6 @@ def f(m, y_true, y_pred):
             ]
         )
 
-        metric0 = metric
         metric = f(metric, y_true, y_pred)
         assert N == 1
         assert metric.compute() == 0.8

tests/metrics/test_metrics.py

@@ -0,0 +1,64 @@
+import jax
+import jax.numpy as jnp
+import pytest
+
+import treex as tx
+from treex import metrics
+
+
+class TestAccuracy:
+    def test_list(self):
+
+        N = 0
+
+        @jax.jit
+        def f(m, y_true, y_pred):
+            nonlocal N
+            N += 1
+            m(y_true=y_true, y_pred=y_pred)
+            return m
+
+        metrics = tx.metrics.Metrics(
+            [
+                tx.metrics.Accuracy(num_classes=10),
+                tx.metrics.Accuracy(num_classes=10),
+            ]
+        )
+        y_true = jnp.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])[None, None, None, :]
+        y_pred = jnp.array([0, 1, 2, 3, 0, 5, 6, 7, 0, 9])[None, None, None, :]
+
+        metrics = f(metrics, y_true, y_pred)
+        assert N == 1
+        assert metrics.compute() == {"accuracy": 0.8, "accuracy2": 0.8}
+
+        metrics = f(metrics, y_true, y_pred)
+        assert N == 1
+        assert metrics.compute() == {"accuracy": 0.8, "accuracy2": 0.8}
+
+    def test_dict(self):
+
+        N = 0
+
+        @jax.jit
+        def f(m, y_true, y_pred):
+            nonlocal N
+            N += 1
+            m(y_true=y_true, y_pred=y_pred)
+            return m
+
+        metrics = tx.metrics.Metrics(
+            dict(
+                a=tx.metrics.Accuracy(num_classes=10),
+                b=tx.metrics.Accuracy(num_classes=10),
+            )
+        )
+        y_true = jnp.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])[None, None, None, :]
+        y_pred = jnp.array([0, 1, 2, 3, 0, 5, 6, 7, 0, 9])[None, None, None, :]
+
+        metrics = f(metrics, y_true, y_pred)
+        assert N == 1
+        assert metrics.compute() == {"a/accuracy": 0.8, "b/accuracy": 0.8}
+
+        metrics = f(metrics, y_true, y_pred)
+        assert N == 1
+        assert metrics.compute() == {"a/accuracy": 0.8, "b/accuracy": 0.8}

treex/metrics/__init__.py

@@ -1,8 +1,5 @@
-import logging as __logging
-
-_logger = __logging.getLogger("treex")
-_logger.addHandler(__logging.StreamHandler())
-_logger.setLevel(__logging.INFO)
-
-from .classification.accuracy import Accuracy
+from .accuracy import Accuracy
+from .mean import Mean
 from .metric import Metric
+from .metrics import Metrics
+from .reduce import Reduce, Reduction