[FEA] New WholeGraph Feature Store for PyG #4432
Conversation
…to bug_post_processing
…to bug_post_processing
…to bug_post_processing
…to bug_post_processing
…to bug_post_processing
…to bug_post_processing
…to bug_post_processing
There was a problem hiding this comment.
Have reviewed the following files as requested by @alexbarghi-nv . Have requested changes mostly around the examples . The feature_store and test look good to me.
python/cugraph-pyg/cugraph_pyg/data/feature_store.py
python/cugraph-pyg/cugraph_pyg/tests/data/test_feature_store_mg.py
python/cugraph-pyg/cugraph_pyg/examples/gcn_dist_sg.py
| wall_clock_start = time.perf_counter() | ||
| device = torch.device("cuda") | ||
|
|
||
| from ogb.nodeproppred import PygNodePropPredDataset # noqa: E402 |
There was a problem hiding this comment.
Suggest fixing import order into Separate Standard, Third-party, and Local Imports.
| dataset = PygNodePropPredDataset(name=args.dataset, root=args.dataset_root) | ||
| split_idx = dataset.get_idx_split() | ||
| data = dataset[0] | ||
|
|
||
| graph_store = cugraph_pyg.data.GraphStore() | ||
| graph_store[ | ||
| ("node", "rel", "node"), "coo", False, (data.num_nodes, data.num_nodes) | ||
| ] = data.edge_index | ||
|
|
||
| feature_store = cugraph_pyg.data.TensorDictFeatureStore() | ||
| feature_store["node", "x"] = data.x | ||
| feature_store["node", "y"] = data.y |
There was a problem hiding this comment.
Should have function around this to keep it cleaner and not have lingering variables like dataset lying around as they consume memory. Probably return the num_features instead of having dataset live throughput the program duration.
| offset = sizes[:rank].sum() if rank > 0 else 0 | ||
|
|
||
| wg_embedding.scatter( | ||
| tensor.clone(memory_format=torch.contiguous_format).cuda(), |
There was a problem hiding this comment.
Can we use tensor.cuda().contiguous() here since .clone() creates a separate copy.?
There was a problem hiding this comment.
Tried that; it does not work
There was a problem hiding this comment.
Was the tensor on the same device before the clone() call? Maybe we can add a fixme here to note the workaround if we didn't meant to create a deepcopy.
There was a problem hiding this comment.
The tensor was on CPU prior to the call
There was a problem hiding this comment.
I'll make an issue
|
/merge |
Reimplements the WG feature store for PyG using the
FeatureStoreinterface.Merge after #4384
Closes rapidsai/wholegraph#47
Closes #4399