Add categorisation to all wandb metrics (#136)

Author: alan-cooney

sparse_autoencoder/loss/abstract_loss.py

@@ -135,7 +135,7 @@ def batch_scalar_loss_with_log(
             )
 
         # Add in the current loss module's metric
-        log_name = self.log_name()
+        log_name = "train/loss/" + self.log_name()
         metrics[log_name] = current_module_loss.detach().cpu().item()
 
         return current_module_loss, metrics

sparse_autoencoder/loss/decoded_activations_l2.py

@@ -27,7 +27,7 @@ class L2ReconstructionLoss(AbstractLoss):
         >>> unused_activations = torch.zeros_like(input_activations)
         >>> # Outputs both loss and metrics to log
         >>> loss(input_activations, unused_activations, output_activations)
-        (tensor(5.5000), {'l2_reconstruction_loss': 5.5})
+        (tensor(5.5000), {'train/loss/l2_reconstruction_loss': 5.5})
     """
 
     _reduction: LossReductionType

sparse_autoencoder/loss/learned_activations_l1.py

@@ -124,8 +124,8 @@ def batch_scalar_loss_with_log(
                 batch_scalar_loss_penalty = absolute_loss_penalty.sum().squeeze()
 
         metrics = {
-            "learned_activations_l1_loss": batch_scalar_loss.item(),
-            self.log_name(): batch_scalar_loss_penalty.item(),
+            "train/loss/" + "learned_activations_l1_loss": batch_scalar_loss.item(),
+            "train/loss/" + self.log_name(): batch_scalar_loss_penalty.item(),
         }
 
         return batch_scalar_loss_penalty, metrics

sparse_autoencoder/loss/tests/test_abstract_loss.py

@@ -64,12 +64,12 @@ def test_batch_scalar_loss_with_log(dummy_loss: DummyLoss) -> None:
         source_activations, learned_activations, decoded_activations
     )
     expected = 2.0  # Mean of [1.0, 2.0, 3.0]
-    assert log["dummy"] == expected
+    assert log["train/loss/dummy"] == expected
 
 
 def test_call_method(dummy_loss: DummyLoss) -> None:
     """Test the call method."""
     source_activations = learned_activations = decoded_activations = torch.ones((1, 3))
     _loss, log = dummy_loss(source_activations, learned_activations, decoded_activations)
     expected = 2.0  # Mean of [1.0, 2.0, 3.0]
-    assert log["dummy"] == expected
+    assert log["train/loss/dummy"] == expected

sparse_autoencoder/metrics/train/capacity.py

@@ -85,5 +85,5 @@ def calculate(self, data: TrainMetricData) -> dict[str, Any]:
         train_batch_capacities_histogram = self.wandb_capacities_histogram(train_batch_capacities)
 
         return {
-            "train_batch_capacities_histogram": train_batch_capacities_histogram,
+            "train/batch_capacities_histogram": train_batch_capacities_histogram,
         }

sparse_autoencoder/metrics/train/feature_density.py

@@ -111,5 +111,5 @@ def calculate(self, data: TrainMetricData) -> dict[str, Any]:
         )
 
         return {
-            "train_batch_feature_density_histogram": train_batch_feature_density_histogram,
+            "train/batch_feature_density_histogram": train_batch_feature_density_histogram,
         }

sparse_autoencoder/metrics/train/l0_norm_metric.py

@@ -22,4 +22,4 @@ def calculate(self, data: TrainMetricData) -> dict[str, float]:
         batch_size = data.learned_activations.size(0)
         n_non_zero_activations = torch.count_nonzero(data.learned_activations)
         batch_average = n_non_zero_activations / batch_size
-        return {"learned_activations_l0_norm": batch_average.item()}
+        return {"train/learned_activations_l0_norm": batch_average.item()}

sparse_autoencoder/metrics/train/tests/test_capacities.py

@@ -62,4 +62,4 @@ def test_calculate_returns_histogram() -> None:
             decoded_activations=activations,
         )
     )
-    assert "train_batch_capacities_histogram" in res
+    assert "train/batch_capacities_histogram" in res

sparse_autoencoder/metrics/train/tests/test_feature_density.py

@@ -41,4 +41,4 @@ def test_calculate_aggregates() -> None:
     )
 
     # Check both metrics are in the result
-    assert "train_batch_feature_density_histogram" in res
+    assert "train/batch_feature_density_histogram" in res

sparse_autoencoder/metrics/train/tests/test_l0_norm_metric.py

@@ -16,4 +16,4 @@ def test_l0_norm_metric() -> None:
     )
     log = l0_norm_metric.calculate(data)
     expected = 3 / 2
-    assert log["learned_activations_l0_norm"] == expected
+    assert log["train/learned_activations_l0_norm"] == expected