Sparsify Dedup Hook - #60
Conversation
| nbr_nids_idx = torch.full_like(hop_nids, -1) | ||
| nbr_nids_idx[hop_mask] = nid_to_idx[hop_nids[hop_mask]] | ||
| batch.nbr_nids_idx.append(nbr_nids_idx) # type: ignore | ||
| batch.global_to_local = lambda x: torch.searchsorted(unique_nids, x) # type: ignore |
There was a problem hiding this comment.
Caching is possible but needs to be clearly useful (right now it's not really) to justify
|
are we getting the sparse tensor implementation speedup from using it on GPU? or it is just the difference between python dictionary and tensor implementation? It is quite interesting if both are on CPU as I thought python dictionaries would be quite efficient for this. Good optimization to have over all. |
shenyangHuang
left a comment
shenyangHuang
left a comment
There was a problem hiding this comment.
looks good, global to local is more intuitive, though it is long to type, but we can keep it for now.
| device = batch.src.device | ||
| z = torch.zeros(len(batch.unique_nids), self.embed_dim, device=device) | ||
|
|
||
| for hop in reversed(range(self.num_layers)): |
There was a problem hiding this comment.
yeah this should go into nn layers, also after we changed the torch.zeros as hard coded feat
| edge_weight = batch.edge_weight if hasattr(batch, 'edge_weight') else None # type: ignore | ||
| z, h_0, c_0 = self.encoder(node_feat, edge_index, edge_weight, h_0, c_0) | ||
| z_src, z_dst, z_neg = z[batch.src_idx], z[batch.dst_idx], z[batch.neg_idx] # type: ignore | ||
| z_src = z[batch.global_to_local(batch.src)] |
There was a problem hiding this comment.
this is a better name. Flagging this to show to users in documentation
| (*batch.nids[hop].shape, self.embed_dim), device=device | ||
| ) | ||
| node_time_feat = self.time_encoder(torch.zeros_like(batch.nids[hop])) | ||
| node_feat = torch.zeros((*seed_nodes.shape, self.embed_dim), device=device) |
There was a problem hiding this comment.
yes, flag to read static node features, even if it is zeros, we should read it
| # If next next hops embeddings exist, use them instead of raw features | ||
| if hop < self.num_layers - 1: | ||
| nbr_feat = z[batch.nbr_nids_idx[hop]] | ||
| nbr_feat = z[batch.global_to_local(nbrs)] |
There was a problem hiding this comment.
global_to_local is very intuitive, though slightly long to type given how often we use it. maybe g2l? not sure
Codecov ReportAll modified and coverable lines are covered by tests ✅ 📢 Thoughts on this report? Let us know! |
2 participants
Purpose
The purpose of this PR is to make an implementation change in our deduplication hook to speed up it's execution.
Given that each batch has ~600 nodes, but the node easily on the order of 10-100k, the following is inefficient in both time and space:
tgm/tgm/hooks.py
Lines 128 to 132 in 5b4949e
Hence, I switched to sparse tensor implementation here ( I tried native Python dict and only got ~5% setup).
Outcome
End to end latency reduction by ~33%
Case (
master)Control (
perf/dict_idx)Key Changes
z[batch.global_to_local(batch.src)]nid_to_idx->global_to_localto make it more understandableRelevant Prs
Close #59