From deed160ad2245d2ba13456f48e5976bca7f6460b Mon Sep 17 00:00:00 2001 From: Simon Meierhans Date: Mon, 26 Oct 2026 02:27:07 -0700 Subject: [PATCH] Accumulate node and link prediction evaluation metrics on device. PiperOrigin-RevId: 1007121121 --- dgf/src/learning/ten_lines/BUILD | 1 + dgf/src/learning/ten_lines/evaluation.py | 20 ++- dgf/src/learning/ten_lines/evaluation_ext.cc | 3 + dgf/src/learning/ten_lines/evaluation_test.py | 37 +++++ .../ten_lines/link_prediction_model.py | 117 ++++++++++++---- .../ten_lines/link_prediction_test.py | 131 ++++++++++++++++++ .../ten_lines/node_prediction_model.py | 85 ++++++++++-- .../ten_lines/node_prediction_test.py | 18 +++ dgf/src/util/util.py | 26 ++++ dgf/src/util/util_test.py | 31 +++++ 10 files changed, 427 insertions(+), 42 deletions(-) diff --git a/dgf/src/learning/ten_lines/BUILD b/dgf/src/learning/ten_lines/BUILD index 6a9f4a7..7822af3 100644 --- a/dgf/src/learning/ten_lines/BUILD +++ b/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", diff --git a/dgf/src/learning/ten_lines/evaluation.py b/dgf/src/learning/ten_lines/evaluation.py index 72aa7bb..e2d256f 100644 --- a/dgf/src/learning/ten_lines/evaluation.py +++ b/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,11 +109,21 @@ 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) @@ -120,13 +131,13 @@ def compute_ranking_metrics( 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) @@ -134,7 +145,8 @@ def hit_at_k(k): 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)) diff --git a/dgf/src/learning/ten_lines/evaluation_ext.cc b/dgf/src/learning/ten_lines/evaluation_ext.cc index 416e351..68c35b8 100644 --- a/dgf/src/learning/ten_lines/evaluation_ext.cc +++ b/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_) { diff --git a/dgf/src/learning/ten_lines/evaluation_test.py b/dgf/src/learning/ten_lines/evaluation_test.py index 7c37a2b..e70028c 100644 --- a/dgf/src/learning/ten_lines/evaluation_test.py +++ b/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 diff --git a/dgf/src/learning/ten_lines/link_prediction_model.py b/dgf/src/learning/ten_lines/link_prediction_model.py index efefa92..8be9474 100644 --- a/dgf/src/learning/ten_lines/link_prediction_model.py +++ b/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, diff --git a/dgf/src/learning/ten_lines/link_prediction_test.py b/dgf/src/learning/ten_lines/link_prediction_test.py index 7b7b446..5a16802 100644 --- a/dgf/src/learning/ten_lines/link_prediction_test.py +++ b/dgf/src/learning/ten_lines/link_prediction_test.py @@ -27,6 +27,7 @@ from dgf.src.io import tf as tf_io from dgf.src.learning.jax.layers import standard from dgf.src.learning.ten_lines import common as common_lib +from dgf.src.learning.ten_lines import evaluation from dgf.src.learning.ten_lines import link_prediction_model from dgf.src.learning.ten_lines import link_prediction_train from dgf.src.sampling import in_memory_sampler as in_memory_sampler_lib @@ -370,6 +371,136 @@ def test_separate_positives_and_negatives(self): test_util.assert_are_equal(self, pos_probs, expected_pos_probs) test_util.assert_are_equal(self, neg_probs, expected_neg_probs) + def test_sum_ranking_metrics(self): + probs = np.array([0.9, 0.1, 0.2, 0.3, 0.5, 0.4], dtype=np.float32) + examples_per_seed_edge = 3 + + metric_sums = link_prediction_model._sum_ranking_metrics( + jnp.asarray(probs), examples_per_seed_edge=examples_per_seed_edge + ) + + pos_probs, neg_probs = ( + link_prediction_model._separate_positives_and_negatives( + probs, examples_per_seed_edge + ) + ) + expected_metrics = evaluation.compute_ranking_metrics(pos_probs, neg_probs) + self.assertSetEqual(set(metric_sums), set(expected_metrics)) + for key, expected_value in expected_metrics.items(): + self.assertIsInstance(metric_sums[key], jax.Array) + self.assertAlmostEqual( + float(metric_sums[key]), float(expected_value) * 2, places=5 + ) + + def test_sum_ranking_metrics_ties(self): + # The second group has ties between the positive and the negatives. + probs = np.array( + [0.9, 0.1, 0.9, 0.3, 0.5, 0.5, 0.5, 0.5], dtype=np.float32 + ) + examples_per_seed_edge = 4 + + metric_sums = link_prediction_model._sum_ranking_metrics( + jnp.asarray(probs), examples_per_seed_edge=examples_per_seed_edge + ) + + pos_probs, neg_probs = ( + link_prediction_model._separate_positives_and_negatives( + probs, examples_per_seed_edge + ) + ) + expected_metrics = evaluation.compute_ranking_metrics(pos_probs, neg_probs) + for key, expected_value in expected_metrics.items(): + self.assertAlmostEqual( + float(metric_sums[key]), float(expected_value) * 2, places=5 + ) + + def test_sum_ranking_metrics_nans(self): + probs = np.array([0.9, 0.1, 0.2, np.nan, 0.5, 0.4], dtype=np.float32) + + metric_sums = link_prediction_model._sum_ranking_metrics( + jnp.asarray(probs), examples_per_seed_edge=3 + ) + + for value in metric_sums.values(): + self.assertTrue(np.isnan(float(value))) + + @parameterized.named_parameters( + dict(testcase_name="host", to_host=True, expected_type=np.ndarray), + dict(testcase_name="device", to_host=False, expected_type=jax.Array), + ) + def test_predict_batch_array_type(self, to_host: bool, expected_type: type): + batches = list( + self.model.predict_batch( + self.graph, + source_node_idxs=[0, 1], + target_node_idxs=[3, 4], + all_combinations=False, + verbose=0, + to_host=to_host, + ) + ) + self.assertNotEmpty(batches) + for batch in batches: + self.assertIsInstance(batch.predictions, expected_type) + + def test_evaluate_with_groups_across_batches(self): + # The batch sizes are not multiples of the group size (1 + 5 negatives), + # so the groups are split across batches. The predictions are rounded to + # create ties. + num_negative_nodes = 5 + examples_per_seed_edge = 1 + num_negative_nodes + batch_sizes = [4, 7, 1, 13] + rng = np.random.default_rng(1) + all_probs = [] + + def fake_predict_batch(graph, source_node_idxs, target_node_idxs, **kwargs): + del graph, kwargs + probs = np.round(rng.uniform(size=len(source_node_idxs)), 1).astype( + np.float32 + ) + all_probs.append(probs) + begin = 0 + batch_idx = 0 + while begin < len(probs): + end = min(begin + batch_sizes[batch_idx % len(batch_sizes)], len(probs)) + yield link_prediction_model.BatchPrediction( + batch_source_node_idxs=np.asarray(source_node_idxs[begin:end]), + batch_target_node_idxs=np.asarray(target_node_idxs[begin:end]), + predictions=jnp.asarray(probs[begin:end]), + ) + begin = end + batch_idx += 1 + + with unittest.mock.patch.object( + self.model, "predict_batch", side_effect=fake_predict_batch + ): + result = self.model.evaluate( + self.graph, + num_eval_steps=50, + num_negative_nodes=num_negative_nodes, + random_seed=1, + verbose=0, + ) + + self.assertLen(all_probs, 1) + pos_probs, neg_probs = ( + link_prediction_model._separate_positives_and_negatives( + all_probs[0], examples_per_seed_edge + ) + ) + self.assertLen(pos_probs, 50) + expected_metrics = evaluation.compute_ranking_metrics(pos_probs, neg_probs) + expected = evaluation.Evaluation( + num_examples=50, + mrr=float(expected_metrics["mrr"]), + auc=float(expected_metrics["auc"]), + hit_at={ + 1: float(expected_metrics["hit_at_1"]), + 5: float(expected_metrics["hit_at_5"]), + }, + ) + test_util.assert_are_equal(self, result, expected, abs_tol=1e-5) + def test_model(self): self.assertEqual(self.model.data().task.target_edgeset, "A_to_B") diff --git a/dgf/src/learning/ten_lines/node_prediction_model.py b/dgf/src/learning/ten_lines/node_prediction_model.py index 52ccfa5..5ea0351 100644 --- a/dgf/src/learning/ten_lines/node_prediction_model.py +++ b/dgf/src/learning/ten_lines/node_prediction_model.py @@ -79,13 +79,14 @@ class BatchPrediction: merged_seed_node_idxs: The indices of the seed nodes in the `normalized_merged_graph`. predictions: The model's output predictions (e.g., probabilities) for the - batch. + batch. A host `np.ndarray` if `predict_batch` is called with + `to_host=True` (default), and a device `jax.Array` otherwise. """ batch_seed_node_idxs: np.ndarray | None normalized_merged_graph: in_memory_graph.InMemoryGraph merged_seed_node_idxs: np.ndarray - predictions: np.ndarray + predictions: np.ndarray | jax.Array # TODO(gbm): Populate. @@ -373,8 +374,24 @@ def predict_batch( seed_node_idxs: common.SeedNodeIdxs, input_features_only: bool, verbose: int = 2, + *, + to_host: bool = True, ) -> Iterator[BatchPrediction]: - """Generate batches of predictions.""" + """Generate batches of predictions. + + Args: + graph: The input graph. + seed_node_idxs: The indices of the seed nodes in the target nodeset. + input_features_only: If true, only the input features (i.e., not the + label) are sampled. + verbose: The verbosity level. + 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. + """ sampling_plan = self._data.sampling_plan if sampling_plan is None: raise ValueError(textwrap.dedent("""\ @@ -434,7 +451,9 @@ def predict_batch( for normalized_batch in util.prefetch_iterator( normalized_batches, prefetch_size=_PREDICT_PREFETCH_SIZE ): - yield self._predict_normalized_batch(live, normalized_batch) + yield self._predict_normalized_batch( + live, normalized_batch, to_host=to_host + ) def _should_cache_normalized_features( self, graph: in_memory_graph.InMemoryGraph, num_seed_nodes: int @@ -463,10 +482,24 @@ def _predict_normalized_batch( self, live: ModelLiveResource, normalized_batch: inference.NormalizedBatch, + *, + to_host: bool = True, ) -> BatchPrediction: - """Predicts on an already merged and normalized (sub-)batch.""" - probabilities = live.apply_core_model(normalized_batch.jax_batch) - predictions = np.asarray(probabilities) + """Predicts on an already merged and normalized (sub-)batch. + + Args: + live: The live model resources. + normalized_batch: The merged and normalized batch. + to_host: If true, the predictions are transferred to the host as + `np.ndarray` (which blocks until the device computation is done). If + false, the predictions stay on device as `jax.Array`. + + Returns: + The batch prediction. + """ + predictions = live.apply_core_model(normalized_batch.jax_batch) + if to_host: + predictions = np.asarray(predictions) return BatchPrediction( batch_seed_node_idxs=normalized_batch.batch_seed_node_idxs, normalized_merged_graph=normalized_batch.normalized_merged_graph, @@ -501,7 +534,9 @@ def predict_on_graph_sample_batch( normalizer=live.normalizer, target_nodeset=self._data.task.target_nodeset, ) - return self._predict_normalized_batch(live, normalized_batch).predictions + return np.asarray( + self._predict_normalized_batch(live, normalized_batch).predictions + ) def evaluate( self, @@ -563,7 +598,11 @@ def evaluate( return self._evaluate_on_batch_iterator( self.predict_batch( - graph, seed_node_idxs, verbose=verbose, input_features_only=False + graph, + seed_node_idxs, + input_features_only=False, + verbose=verbose, + to_host=False, ) ) @@ -652,7 +691,9 @@ def batch_prediction_generator() -> Iterator[BatchPrediction]: verbose=verbose, ) for normalized_batch in normalized_batches: - yield self._predict_normalized_batch(live, normalized_batch) + yield self._predict_normalized_batch( + live, normalized_batch, to_host=False + ) if verbose >= 1: log.info("Evaluating model on generator") @@ -662,6 +703,20 @@ def _evaluate_on_batch_iterator( self, batch_iterator: Iterator[BatchPrediction], ) -> evaluation.Evaluation: + """Evaluates the model on batches of predictions. + + The metrics are accumulated on device, and only transferred to the host at + the end, to avoid a device synchronization at each batch. The only + exception is the classification per-class metrics (e.g., ROC AUC), which + are accumulated on the host. + + Args: + batch_iterator: The batches of predictions, ideally with predictions on + device. + + Returns: + The evaluation. + """ num_good_predictions = jnp.array(0) total_squared_error = jnp.array(0.0) sum_labels = jnp.array(0.0) @@ -677,7 +732,12 @@ def _evaluate_on_batch_iterator( num_classes = self.num_label_classes() accumulator = evaluation.ClassificationEvaluationAccumulator(num_classes) - for batch in batch_iterator: + # Each batch is processed only after the next batch is fetched, i.e., after + # the model computation of the next batch is dispatched to the device. This + # way, the host-side work on the current batch (notably the per-class + # metrics accumulator) overlaps with the device computation of the next + # batch, instead of leaving the device idle. + for batch in util.lookahead_iterator(batch_iterator): labels = batch.normalized_merged_graph.node_sets[target_nodeset].features[ target_normalized_feature ][batch.merged_seed_node_idxs] @@ -692,6 +752,9 @@ def _evaluate_on_batch_iterator( predictions = jnp.argmax(batch.predictions, axis=-1) num_good_predictions += jnp.sum(labels == predictions) assert accumulator is not None + # Note: The per-class metrics accumulator runs on the host, so this + # transfers the predictions to the host. Thanks to the lookahead above, + # the device computes the next batch in the meantime. accumulator.add_predictions( np.asarray(batch.predictions, dtype=np.float32), np.asarray(labels, dtype=np.int32), diff --git a/dgf/src/learning/ten_lines/node_prediction_test.py b/dgf/src/learning/ten_lines/node_prediction_test.py index 5920618..c3503fe 100644 --- a/dgf/src/learning/ten_lines/node_prediction_test.py +++ b/dgf/src/learning/ten_lines/node_prediction_test.py @@ -543,6 +543,24 @@ def test_predict_batch_splits_oversized_batches(self): seed_node_idxs, ) + @parameterized.named_parameters( + dict(testcase_name="host", to_host=True, expected_type=np.ndarray), + dict(testcase_name="device", to_host=False, expected_type=jax.Array), + ) + def test_predict_batch_array_type(self, to_host: bool, expected_type: type): + batches = list( + self.model.predict_batch( + graph=self.graph, + seed_node_idxs=[0, 1, 2], + input_features_only=True, + verbose=0, + to_host=to_host, + ) + ) + self.assertNotEmpty(batches) + for batch in batches: + self.assertIsInstance(batch.predictions, expected_type) + def test_predict_without_labels(self): graph_without_labels = copy.deepcopy(self.graph) del graph_without_labels.node_sets["client"].features["categorical_label"] diff --git a/dgf/src/util/util.py b/dgf/src/util/util.py index 598d315..8ff78bb 100644 --- a/dgf/src/util/util.py +++ b/dgf/src/util/util.py @@ -14,6 +14,7 @@ """Utility functions for the DGF project.""" +import collections from collections.abc import Iterator import contextlib import math @@ -427,3 +428,28 @@ def prefetch_iterator( if prefetch_size <= 0: return iterator return BackgroundPrefetcher(iterator, prefetch_size=prefetch_size) + + +def lookahead_iterator(iterator: Iterator[T], lookahead: int = 1) -> Iterator[T]: + """Yields each element only after the next `lookahead` elements are fetched. + + Unlike `prefetch_iterator`, no background thread is used: the next elements + are fetched on the calling thread, right before the current element is + yielded. This is useful when fetching an element starts some asynchronous + work (e.g., dispatching a JAX computation to the device): the asynchronous + work of the next elements then overlaps with the processing of the current + element by the caller. + + Args: + iterator: The source iterator. + lookahead: The number of elements to fetch ahead of the yielded element. + + Yields: + The elements of `iterator`, in order. + """ + buffer = collections.deque() + for item in iterator: + buffer.append(item) + if len(buffer) > lookahead: + yield buffer.popleft() + yield from buffer diff --git a/dgf/src/util/util_test.py b/dgf/src/util/util_test.py index 9b798df..1f8fe63 100644 --- a/dgf/src/util/util_test.py +++ b/dgf/src/util/util_test.py @@ -379,5 +379,36 @@ def failing_iterator(): next(it) +class LookaheadIteratorTest(absltest.TestCase): + + def test_lookahead_basic(self): + items = list(range(10)) + self.assertEqual(list(util.lookahead_iterator(iter(items))), items) + self.assertEqual( + list(util.lookahead_iterator(iter(items), lookahead=3)), items + ) + + def test_lookahead_short_and_empty(self): + self.assertEqual(list(util.lookahead_iterator(iter([1]), lookahead=3)), [1]) + self.assertEqual(list(util.lookahead_iterator(iter([]))), []) + + def test_lookahead_fetches_next_before_yield(self): + num_fetched = 0 + + def counting_iterator(): + nonlocal num_fetched + for i in range(4): + num_fetched += 1 + yield i + + it = util.lookahead_iterator(counting_iterator(), lookahead=1) + self.assertEqual(next(it), 0) + self.assertEqual(num_fetched, 2) + self.assertEqual(next(it), 1) + self.assertEqual(num_fetched, 3) + self.assertEqual(list(it), [2, 3]) + self.assertEqual(num_fetched, 4) + + if __name__ == "__main__": absltest.main()