Separate the metrics variables from non-trainable variables.

keras_core/layers/layer.py

@@ -514,10 +514,13 @@ def trainable(self, value):
 
     @property
     def variables(self):
-        """List of all layer state, including metric variables and random seeds.
+        """List of all layer state, including random seeds.
 
         This extends `layer.weights` to include all state used by the layer
-        including state for metrics and `SeedGenerator`s.
+        including `SeedGenerator`s.
+
+        Note that metrics variables are not included here, use
+        `metrics_variables` to visit all the metric variables.
         """
         # Return all `Variables` associate with the layer including metrics
         # and random seeds. Also deduplicate them.
@@ -527,8 +530,6 @@ def variables(self):
             if id(v) not in seen_ids:
                 variables.append(v)
                 seen_ids.add(id(v))
-        for m in self._metrics:
-            variables.extend(m.variables)
         for sg in self._seed_generators:
             variables.append(sg.state)
         for layer in self._layers:
@@ -602,6 +603,17 @@ def non_trainable_weights(self):
             return self.weights
         return [v for v in self.weights if not v.trainable]
 
+    @property
+    def metrics_variables(self):
+        """List of all metric variables."""
+        vars = []
+        for metric in self._metrics:
+            vars.extend(metric.variables)
+        for layer in self._layers:
+            for metric in layer._metrics:
+                vars.extend(metric.variables)
+        return vars
+
     def get_weights(self):
         """Return the values of `layer.weights` as a list of NumPy arrays."""
         return [v.numpy() for v in self.weights]

keras_core/layers/layer_test.py

@@ -3,6 +3,7 @@
 
 from keras_core import backend
 from keras_core import layers
+from keras_core import metrics
 from keras_core import models
 from keras_core import ops
 from keras_core import testing
@@ -214,6 +215,45 @@ def call(self, x):
         self.assertLen(layer.inner_layer.weights, 8)
         self.assertLen(layer.weights, 8)
 
+    def test_metric_tracking(self):
+        class LayerWithMetric(layers.Layer):
+            def __init__(self, units):
+                super().__init__()
+                self.dense = layers.Dense(units)
+                self.metric = metrics.MeanSquaredError(name="my_metric")
+
+            def build(self, input_shape):
+                self.dense.build(input_shape)
+
+            def call(self, x):
+                return self.dense(x)
+
+        class NestedLayerWithMetric(layers.Layer):
+            def __init__(self, units):
+                super().__init__()
+                self.layer_with_metric = LayerWithMetric(units)
+                self.metric = metrics.MeanSquaredError(name="my_metric")
+
+            def build(self, input_shape):
+                self.layer_with_metric.build(input_shape)
+
+            def call(self, x):
+                return self.layer_with_metric(x)
+
+        layer = LayerWithMetric(3)
+        layer.build((1, 3))
+
+        self.assertLen(layer.metrics_variables, 2)
+        self.assertLen(layer.trainable_variables, 2)
+        self.assertLen(layer.non_trainable_variables, 0)
+
+        layer = NestedLayerWithMetric(3)
+        layer.build((1, 3))
+
+        self.assertLen(layer.metrics_variables, 4)
+        self.assertLen(layer.trainable_variables, 2)
+        self.assertLen(layer.non_trainable_variables, 0)
+
     def test_build_on_call(self):
         class LayerWithUnbuiltState(layers.Layer):
             def __init__(self, units):

keras_core/trainers/trainer_test.py

@@ -29,7 +29,7 @@
 
 
 # A model is just a layer mixed in with a Trainer.
-class ExampleModel(layers.Dense, Trainer):
+class ExampleModel(Trainer, layers.Dense):
     def __init__(self, units):
         layers.Dense.__init__(
             self,
@@ -40,7 +40,7 @@ def __init__(self, units):
         Trainer.__init__(self)
 
 
-class StructModel(layers.Layer, Trainer):
+class StructModel(Trainer, layers.Layer):
     def __init__(self, units):
         layers.Layer.__init__(self)
         Trainer.__init__(self)
@@ -62,7 +62,7 @@ def call(self, x):
         }
 
 
-class ListModel(layers.Layer, Trainer):
+class ListModel(Trainer, layers.Layer):
     def __init__(self, units):
         layers.Layer.__init__(self)
         Trainer.__init__(self)
@@ -82,7 +82,7 @@ def call(self, x):
         return self.dense_1(x[0]) + self.dense_2(x[1])
 
 
-class TrainingTestingLayer(layers.Layer, Trainer):
+class TrainingTestingLayer(Trainer, layers.Layer):
     def __init__(self, **kwargs):
         layers.Layer.__init__(self, **kwargs)
         Trainer.__init__(self)
@@ -96,7 +96,7 @@ def call(self, x, training=False):
 class TestTrainer(testing.TestCase, parameterized.TestCase):
     @pytest.mark.requires_trainable_backend
     def test_metric_tracking(self):
-        class ModelWithMetric(layers.Dense, Trainer):
+        class ModelWithMetric(Trainer, layers.Dense):
             def __init__(self, units):
                 layers.Dense.__init__(
                     self,
@@ -132,6 +132,7 @@ def __init__(self, units):
 
         # And those weights are tracked at the model level
         self.assertEqual(len(model.metrics_variables), 6)
+        self.assertLen(model.non_trainable_variables, 0)
 
         # Models with only weighted_metrics should have the same 3 metrics
         model_weighted = ModelWithMetric(units=3)
@@ -361,7 +362,7 @@ def test_adds_loss_scaling_optimizer(self):
         reason="half precision unsupported on torch CPU.",
     )
     def test_loss_scaling_prevents_underflow(self):
-        class DeepModel(layers.Layer, Trainer):
+        class DeepModel(Trainer, layers.Layer):
             def __init__(self):
                 layers.Layer.__init__(self, dtype="mixed_float16")
                 Trainer.__init__(self)