Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add an RGCNConv model using the primitives in cugraph-ops #6278

Merged
merged 22 commits into from
Feb 3, 2023
Merged

Add an RGCNConv model using the primitives in cugraph-ops #6278

merged 22 commits into from
Feb 3, 2023

Conversation

tingyu66
Copy link
Contributor

@tingyu66 tingyu66 commented Dec 21, 2022

This PR adds an RGCN model torch_geometric.nn.RGCNConvCuGraph that leverages the optimized aggregation functions in cugraph-ops.

Fixes https://github.com/rapidsai/cugraph-ops/issues/180.

It requires the package pylibcugraphops >= 23.02, which is going to be released by next February. I have tested it using the nightly build.

The user interface largely resembles RGCNConv. Our aggregation function has two modes: full-graph and message flow graph (MFG) mode, with the latter being more performant. That's why our forward function has to have extra arguments like num_nodes, from_neighbor_sampler and max_num_neighbors, to help convert PyG edge_index to cugraph-ops graph type and specify which mode to opt for.

Here are some preliminary performance results using MAG dataset and 2-layer RGCN model, with batch_size=1024, num_neighbors=[25,20] in NeighborLoader. I set feature dim to 16 and did not use the provided paper features, as this is just a synthetic benchmark.

time (s) per epoch RGCNConv with pyg_lib RGCNConvCuGraph
forward 1.66 1.14
backward 2.62 1.52

Right now, since cugraph-ops requires csc/csr representation of the graph, the forward time is dominated by the coo-to-csc conversion using the utility for SparseTensor, taking 50~70% of the forward time. We can expect a higher speedup if 1). dataloader can emit csc graph and 2). dataloader can output a list of MFG, instead of a fused subgraph.

@codecov
Copy link

codecov bot commented Jan 26, 2023

Codecov Report

Merging #6278 (2ee92ba) into master (06c3043) will decrease coverage by 0.01%.
The diff coverage is 84.61%.

❗ Current head 2ee92ba differs from pull request most recent head 0274fe9. Consider uploading reports for the commit 0274fe9 to get more accurate results

@@            Coverage Diff             @@
##           master    #6278      +/-   ##
==========================================
- Coverage   86.25%   86.25%   -0.01%     
==========================================
  Files         417      418       +1     
  Lines       22779    22791      +12     
==========================================
+ Hits        19649    19659      +10     
- Misses       3130     3132       +2     
Impacted Files Coverage Δ
torch_geometric/nn/conv/cugraph/rgcn_conv.py 80.00% <80.00%> (ø)
torch_geometric/nn/conv/__init__.py 100.00% <100.00%> (ø)
torch_geometric/nn/conv/cugraph/__init__.py 100.00% <100.00%> (ø)
torch_geometric/nn/conv/cugraph/base.py 85.71% <100.00%> (ø)

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more

@tingyu66 tingyu66 marked this pull request as ready for review January 26, 2023 23:27
rusty1s added a commit that referenced this pull request Feb 2, 2023
This PR adds a GraphSAGE model `torch_geometric.nn.SAGEConvCuGraph` that
leverages the optimized aggregation function in cugraph-ops, similar to
#6278.

Fixes rapidsai/cugraph-ops#178.

---------

Co-authored-by: Matthias Fey <matthias.fey@tu-dortmund.de>
@tingyu66
Copy link
Contributor Author

tingyu66 commented Feb 2, 2023

@rusty1s I have updated RGCNConv and tested it locally. Feel free to give it a go.

rusty1s added a commit that referenced this pull request Feb 3, 2023
This PR adds a GAT model `torch_geometric.nn.GATConvCuGraph` that
leverages the optimized aggregation function in cugraph-ops, similar to
#6278.

Fixes rapidsai/cugraph-ops#182.

---------

Co-authored-by: Matthias Fey <matthias.fey@tu-dortmund.de>
@rusty1s rusty1s enabled auto-merge (squash) February 3, 2023 19:23
@rusty1s rusty1s merged commit 5a277a2 into pyg-team:master Feb 3, 2023
@tingyu66 tingyu66 deleted the cugraphops_rgcn branch February 7, 2023 14:43
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

2 participants