Vectorize Recency Sampler Using Circular Buffers - #62
Conversation
Jacob-Chmura
changed the title
Codecov ReportAttention: Patch coverage is
📢 Thoughts on this report? Let us know! |
| fake_times = torch.randint( | ||
| int(batch.time.min().item()), | ||
| int(batch.time.max().item()), | ||
| int(batch.time.max().item()) + 1, |
There was a problem hiding this comment.
this is to be inclusive of the maxtime right
There was a problem hiding this comment.
maybe add a comment for us to remember in the future
timestamps for negatives destinations in a batch is sampled from [batch_min_time, batch_max_time] inclusive.
| @@ -293,7 +303,7 @@ def __call__(self, dg: DGraph, batch: DGBatch) -> DGBatch: | |||
| # leakage, making the prediction easier than it should be. | |||
There was a problem hiding this comment.
what does this mean? should we remove the leakage comment?
There was a problem hiding this comment.
This explains why we chose our heuristic for assigning timestamps on negative edges. We don't do random timestamps because this would give an implicit signal to the model about which edges are positive and negative.
| def _get_recency_indices(self, node_ids: torch.Tensor, k: int) -> torch.Tensor: | ||
| ptr = self._nbr_ptr[node_ids].unsqueeze(1) | ||
| offsets = torch.arange(k, device=node_ids.device).unsqueeze(0) | ||
| indices = (ptr - 1 - offsets) % self._max_nbrs |
There was a problem hiding this comment.
looks correct, though with pointer and everything, should make sure it is working as intended. What happens when a node has no neighbors whatsoever?
2 participants
Purpose
The purpose of this PR is to make an implementation change in our recency sampler hook to speed up it's execution.
Specifically, I moved away from python queues and wrote a fully torch-native implementation using circular buffers, which can be stored on GPU.
Outcome (independent of sparse dedup hook improvements #60)
End to end latency reduction by ~81%
Case (
perf/dict_idx)Control (
perf/recency)Key Changes
Relevant Prs
Close #61