Accumulate node and link prediction evaluation metrics on device.

dgf/src/learning/ten_lines/BUILD

@@ -462,6 +462,7 @@ py_test(
     shard_count = 4,
     deps = [
         ":common",
+        ":evaluation",
         ":link_prediction_model",
         ":link_prediction_train",
         #"//learning/brain/research/jax:gpu_support",

dgf/src/learning/ten_lines/evaluation.py

@@ -14,7 +14,7 @@
 
 """Evaluation utilities for ten_lines model."""
 
-from typing import Any
+from typing import Any, Literal
 from dgf.src.data import evaluation as evaluation_data_lib
 from dgf.src.learning.ten_lines import evaluation_ext
 import jax
@@ -100,6 +100,7 @@ def populate_evaluation(self, evaluation: Evaluation) -> None:
 def compute_ranking_metrics(
     pos_scores: jax.Array | np.ndarray,
     neg_scores: jax.Array | np.ndarray,
+    reduction: Literal["mean", "sum"] = "mean",
 ) -> dict[str, jax.Array | np.ndarray]:
   """Computes ranking metrics (MRR, Hit@1, Hit@5, AUC) from scores.
 
@@ -108,33 +109,44 @@ def compute_ranking_metrics(
   Args:
     pos_scores: Scores for positive edges, shape (B,).
     neg_scores: Scores for negative edges, shape (B * N,).
+    reduction: How the per-positive-edge metrics are reduced over the B
+      positive edges. "mean" returns the average metrics. "sum" returns the
+      sum of the metrics, which is useful to accumulate metrics over batches
+      of different sizes without the extra rounding of `mean * B`.
 
   Returns:
     A dictionary containing the computed metrics.
   """
   xp = np if isinstance(pos_scores, np.ndarray) else jnp
+  if reduction == "mean":
+    reduce_fn = xp.mean
+  elif reduction == "sum":
+    reduce_fn = xp.sum
+  else:
+    raise ValueError(f"Unknown reduction: {reduction!r}")
 
   neg_scores_reshaped = neg_scores.reshape(pos_scores.shape[0], -1)
 
   greater = xp.sum(neg_scores_reshaped > pos_scores[:, None], axis=-1)
   equal = xp.sum(neg_scores_reshaped == pos_scores[:, None], axis=-1)
 
   mid_ranks = greater + 1 + equal / 2.0
-  mrr = xp.mean(1.0 / mid_ranks)
+  mrr = reduce_fn(1.0 / mid_ranks)
 
   def hit_at_k(k):
     slots = k - greater
     total_tied = equal + 1
     prob = xp.clip(slots / total_tied, 0.0, 1.0)
-    return xp.mean(prob)
+    return reduce_fn(prob)
 
   hit_at_1 = hit_at_k(1)
   hit_at_5 = hit_at_k(5)
 
   auc_matrix = (pos_scores[:, None] > neg_scores_reshaped).astype(
       xp.float32
   ) + 0.5 * (pos_scores[:, None] == neg_scores_reshaped).astype(xp.float32)
-  auc = xp.mean(auc_matrix)
+  # Per-positive-edge AUC, reduced over the positive edges.
+  auc = reduce_fn(xp.mean(auc_matrix, axis=-1))
 
   invalid = xp.any(xp.isnan(pos_scores)) | xp.any(xp.isnan(neg_scores))
 

dgf/src/learning/ten_lines/evaluation_ext.cc

@@ -56,6 +56,9 @@ class ClassificationEvaluationAccumulator {
           std::to_string(num_examples));
     }
 
+    // The arrays are kept alive by the caller, so the GIL can be released. This
+    // lets other Python threads (e.g., batch prefetching) run in the meantime.
+    nb::gil_scoped_release release;
     for (size_t i = 0; i < num_examples; i++) {
       int32_t true_class = target_view(i);
       if (true_class < 0 || true_class >= num_classes_) {

dgf/src/learning/ten_lines/evaluation_test.py

@@ -82,6 +82,43 @@ def test_compute_ranking_metrics_nans(self):
     self.assertTrue(np.isnan(float(metrics["hit_at_5"])))
     self.assertTrue(np.isnan(float(metrics["auc"])))
 
+  def test_compute_ranking_metrics_sum_reduction(self):
+    # Two positive edges with 4 negatives each, including ties.
+    pos_scores = np.array([0.5, 0.8], dtype=np.float32)
+    neg_scores = np.array(
+        [0.2, 0.5, 0.7, 0.9, 0.8, 0.1, 0.8, 0.3], dtype=np.float32
+    )
+    num_pos = len(pos_scores)
+    for xp in (np, jnp):
+      mean_metrics = evaluation_lib.compute_ranking_metrics(
+          xp.asarray(pos_scores), xp.asarray(neg_scores)
+      )
+      sum_metrics = evaluation_lib.compute_ranking_metrics(
+          xp.asarray(pos_scores), xp.asarray(neg_scores), reduction="sum"
+      )
+      self.assertSetEqual(set(sum_metrics), set(mean_metrics))
+      for key, mean_value in mean_metrics.items():
+        self.assertAlmostEqual(
+            float(sum_metrics[key]), float(mean_value) * num_pos, places=5
+        )
+
+  def test_compute_ranking_metrics_sum_reduction_nans(self):
+    pos_scores = np.array([np.nan, 0.5])
+    neg_scores = np.array([0.2] * 16)
+    metrics = evaluation_lib.compute_ranking_metrics(
+        pos_scores, neg_scores, reduction="sum"
+    )
+    for value in metrics.values():
+      self.assertTrue(np.isnan(float(value)))
+
+  def test_compute_ranking_metrics_unknown_reduction(self):
+    with self.assertRaisesRegex(ValueError, "Unknown reduction"):
+      evaluation_lib.compute_ranking_metrics(
+          np.array([0.5]),
+          np.array([0.2]),
+          reduction="max",  # pyrefly: ignore[bad-argument-type]
+      )
+
   def test_classification_evaluation_accumulator(self):
     np.random.seed(42)
     num_classes = 3

dgf/src/learning/ten_lines/link_prediction_model.py

@@ -18,6 +18,8 @@
 
 from collections.abc import Callable, Iterator
 import dataclasses
+import functools
+import math
 import os
 from typing import Any, Literal
 
@@ -147,11 +149,19 @@ class ModelLiveResource:
 
 @dataclasses.dataclass
 class BatchPrediction:
-  """Prediction yielded by "predict_batch" function."""
+  """Prediction yielded by "predict_batch" function.
+
+  Attributes:
+    batch_source_node_idxs: The indices of the source nodes.
+    batch_target_node_idxs: The indices of the target nodes.
+    predictions: The predicted probabilities. A host `np.ndarray` if
+      `predict_batch` is called with `to_host=True` (default), and a device
+      `jax.Array` otherwise.
+  """
 
   batch_source_node_idxs: np.ndarray
   batch_target_node_idxs: np.ndarray
-  predictions: np.ndarray
+  predictions: np.ndarray | jax.Array
 
 
 def _interleave_positives_and_negatives(
@@ -201,8 +211,8 @@ def _interleave_positives_and_negatives(
 
 
 def _separate_positives_and_negatives(
-    probs: np.ndarray, examples_per_seed_edge: int
-) -> tuple[np.ndarray, np.ndarray]:
+    probs: np.ndarray | jax.Array, examples_per_seed_edge: int
+) -> tuple[np.ndarray | jax.Array, np.ndarray | jax.Array]:
   """Separates positive and negative probabilities from a combined array.
 
   Toy Example:
@@ -229,6 +239,29 @@ def _separate_positives_and_negatives(
   return pos_probs, neg_probs
 
 
+@functools.partial(jax.jit, static_argnames=("examples_per_seed_edge",))
+def _sum_ranking_metrics(
+    probs: jax.Array, examples_per_seed_edge: int
+) -> dict[str, jax.Array]:
+  """Computes the sum (over the seed edges) of the ranking metrics.
+
+  Args:
+    probs: Combined probabilities array, shape (num_groups *
+      examples_per_seed_edge,). See `_separate_positives_and_negatives`.
+    examples_per_seed_edge: Number of examples per seed edge (1 pos + K neg).
+
+  Returns:
+    The sum (over the seed edges) of the per-seed-edge ranking metrics of
+    `evaluation.compute_ranking_metrics`.
+  """
+  pos_probs, neg_probs = _separate_positives_and_negatives(
+      probs, examples_per_seed_edge
+  )
+  return evaluation.compute_ranking_metrics(  # pyrefly: ignore[bad-return]
+      pos_probs, neg_probs, reduction="sum"
+  )
+
+
 def _rename_spec(spec: Any, name: str) -> Any:
   """Returns a copy of a tensor spec with a new name.
 
@@ -577,8 +610,27 @@ def predict_batch(
       verbose: int = 2,
       source_sampler: in_memory_sampler_lib.Sampler | None = None,
       target_sampler: in_memory_sampler_lib.Sampler | None = None,
+      *,
+      to_host: bool = True,
   ) -> Iterator[BatchPrediction]:
-    """Generate batches of predictions."""
+    """Generate batches of predictions.
+
+    Args:
+      graph: The input graph.
+      source_node_idxs: The indices of the source nodes.
+      target_node_idxs: The indices of the target nodes.
+      all_combinations: If true, predict all the source x target combinations.
+      edge_timestamps: Optional timestamps of the edges to predict.
+      verbose: The verbosity level.
+      source_sampler: Optional sampler for the source nodes.
+      target_sampler: Optional sampler for the target nodes.
+      to_host: If true, the predictions are transferred to the host as
+        `np.ndarray`. If false, the predictions stay on device as `jax.Array`,
+        which avoids a device synchronization at each batch.
+
+    Yields:
+      The batches of predictions.
+    """
     live = self._get_live()
 
     np_source_node_idxs = np.asarray(source_node_idxs)
@@ -714,7 +766,9 @@ def merge_and_predict(
       )
 
       logits = live.apply_core_model(batch)
-      predictions = np.asarray(jax.nn.sigmoid(logits))
+      predictions = jax.nn.sigmoid(logits)
+      if to_host:
+        predictions = np.asarray(predictions)
 
       yield BatchPrediction(
           batch_source_node_idxs=sub_src,
@@ -1396,15 +1450,20 @@ def evaluate(
         neighbor_idxs.neg_trg_node_idxs,
     )
 
-    total_mrr = 0.0
-    total_hit_at_1 = 0.0
-    total_hit_at_5 = 0.0
-    total_auc = 0.0
+    # The metrics are summed on device for each batch, and only transferred to
+    # the host at the end, to avoid a device synchronization at each batch. The
+    # per-batch sums are then accumulated on the host in float64, to avoid the
+    # precision loss of a float32 accumulation over many batches.
+    # TODO(simonmeierhans): Bound the memory usage by periodically compacting
+    # `batch_metric_sums` (e.g., every 10k batches, transfer the pending sums to
+    # the host, accumulate them into float64 running totals, and clear the
+    # list). This keeps virtually the same speed (one sync every 10k batches).
+    batch_metric_sums: list[dict[str, jax.Array]] = []
 
     # Each group contains one positive prediction, followed by
     # num_negative_nodes corresponding predictions.
     examples_per_seed_edge = 1 + num_negative_nodes
-    carry_over_probs = np.array([])
+    carry_over_probs = jnp.zeros((0,), dtype=jnp.float32)
 
     full_timestamps = None
     if self._data.temporal_sampling:
@@ -1418,12 +1477,14 @@ def evaluate(
         graph,
         full_src.tolist(),
         full_trg.tolist(),
+        all_combinations=False,
         edge_timestamps=full_timestamps,
         verbose=verbose,
         source_sampler=source_sampler,
         target_sampler=target_sampler,
+        to_host=False,
     ):
-      combined_probs = np.concatenate(
+      combined_probs = jnp.concatenate(
           [carry_over_probs, batch_pred.predictions]
       )
 
@@ -1435,22 +1496,11 @@ def evaluate(
         complete_probs = combined_probs[:complete_part_len]
         carry_over_probs = combined_probs[complete_part_len:]
 
-        pos_probs, neg_probs = _separate_positives_and_negatives(
-            complete_probs, examples_per_seed_edge
-        )
-
-        ranking_metrics = evaluation.compute_ranking_metrics(
-            pos_probs, neg_probs
-        )
-
-        total_mrr += float(ranking_metrics["mrr"]) * num_complete_groups
-        total_hit_at_1 += (
-            float(ranking_metrics["hit_at_1"]) * num_complete_groups
-        )
-        total_hit_at_5 += (
-            float(ranking_metrics["hit_at_5"]) * num_complete_groups
+        batch_metric_sums.append(
+            _sum_ranking_metrics(
+                complete_probs, examples_per_seed_edge=examples_per_seed_edge
+            )
         )
-        total_auc += float(ranking_metrics["auc"]) * num_complete_groups
       else:
         carry_over_probs = combined_probs
 
@@ -1460,6 +1510,19 @@ def evaluate(
           " examples left over. Total number of examples must be a multiple of"
           f" {examples_per_seed_edge}."
       )
+    if not batch_metric_sums:
+      raise ValueError("No examples to evaluate: the seed edges are empty.")
+
+    # Single device-to-host transfer.
+    host_batch_metric_sums = jax.device_get(batch_metric_sums)
+
+    def total(key: str) -> float:
+      return math.fsum(float(m[key]) for m in host_batch_metric_sums)
+
+    total_mrr = total("mrr")
+    total_hit_at_1 = total("hit_at_1")
+    total_hit_at_5 = total("hit_at_5")
+    total_auc = total("auc")
 
     return evaluation.Evaluation(
         num_examples=num_examples,