GPSE Implementation

[semihcanturk]

Graph Positional and Structural Encoder implementation as per #8310. Adapted from the original repository: https://github.com/G-Taxonomy-Workgroup/GPSE. This version is a standalone implementation that is decoupled from GraphGym, and thus aims for better accessibility and a smoother integration into PyG. While the priority of this PR is to enable loading and using pre-trained models in plug-and-play fashion, it also includes the custom loss function used to train the model. Nevertheless, it might be easier to use the original repository for pre-training and fine-tuning new GPSE models for the time being.

This PR includes the following:

• GPSE: The main GPSE module, that generates learned encodings for input graphs.
• Several helper classes (FeatureEncoder, GNNStackStage, IdentityHead, GNNInductiveHybridMultiHead, ResGatedGCNConvGraphGymLayer, Linear, MLP, GeneralMultiLayer, GeneralLayer, BatchNorm1dNode, BatchNorm1dEdge, VirtualNodePatchSingleton) and wrapper functions (GNNPreMP, GNNLayer), all adapted from their GraphGym versions for compatibility and enabling the loading of weights pre-trained using the GraphGym/original version.
• The class method GPSE.from_pretrained() that returns a model with pre-trained weights from the original repository/Zenodo files.
• GPSENodeEncoder, a helper linear/MLP encoder that takes the GPSE encodings precomputed asbatch.pestat_GPSE in the input graphs, maps them to a desired dimension and appends them to node features.
• precompute_GPSE , a function that takes in a GPSE model and a dataset, and precomputes GPSE encodings in-place for a given dataset using the helper function gpse_process_batch.
• The transform AddGPSE, which in similar fashion to AddLaplacianEigenvectorPE and AddRandomWalkPE adds the GPSE encodings to a given graph using the helper function gpse_process
• The testing modules test/test_gpse.py and test/test_add_gpse.py.
• The loss function gpse_loss and helper functions cosim_col_sep and process_batch_idx used in GPSE training.
• A comprehensive example in examples/gpse.py is provided as a separate PR in GPSE example #10118.

This PR has been tested with the latest (25.01-py3) NVIDIA stack, and has works without any issues.

[codecov]

Codecov Report

Attention: Patch coverage is 64.41860% with 153 lines in your changes are missing coverage. Please review.

Project coverage is 89.23%. Comparing base (61c47ee) to head (f24acbc).
Report is 14 commits behind head on master.

❗ Current head f24acbc differs from pull request most recent head f9ce537

Please upload reports for the commit f9ce537 to get more accurate results.

Files	Patch %	Lines
torch_geometric/nn/models/gpse.py	62.68%	153 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #9018      +/-   ##
==========================================
+ Coverage   87.33%   89.23%   +1.90%     
==========================================
  Files         460      472      +12     
  Lines       30385    30594     +209     
==========================================
+ Hits        26536    27301     +765     
+ Misses       3849     3293     -556     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[rusty1s]

Super, let me resolve the linting issue.

[pjspol]

@semihcanturk @rusty1s What is the status on this PyG implementation of GPSE? The GPSE paper says, "For convenience, GPSE has also been integrated into the PyG library to facilitate downstream applications." I was hoping to try it out soon.

[semihcanturk]

Hi @pjspol, and thanks for checking in! The implementation here is ready and is simply awaiting a merge -- I know @rusty1s was unavailable for a while so that has delayed the merge a bit, but I'm expecting it to be merged soon. In the meantime, you should be able to install this branch in your environment to work with GPSE: pip install git+https://github.com/semihcanturk/pytorch_geometric/tree/gpse

[luke-a-thompson]

Hi @semihcanturk. I believe there is a minor ambiguity in the docs for the GPSENodeEncoder. In the expand_x option (line 663), it does not specify if the expansion will be to dim_pe_in or dim_pe_out. I assume it's dim_pe_out, but it still might be worth noting exactly which shape the output will take.

expand_x (bool, optional): Expand node features :obj:`x` from
    :obj:`dim_in` to (:obj:`dim_emb` - :obj:`dim_pe`)

Thanks

[semihcanturk]

Hi @rusty1s, finally got around to this :) There were minor issues re: bn_eps and bn_mom args after your update which are now fixed + fixed docs (thanks for the catch @luke-a-thompson!); I've also updated the branch to latest PyG. Once again I can't pass the linting check for whatever reason, but otherwise we should be good to go here -- lmk if you need anything on my end.

[zzzzzzyc]

What is the status on this PyG implementation of GPSE?I think it will be very helpful for my research!

[puririshi98]

@zzzzzzyc i will work with @semihcanturk to get this merged in but for the meantime you can test this branch in your own work.

please do share your experience with running this
pip uninstall -y torch-geometric; rm -rf pytorch_geometric; git clone -b gpse https://github.com/semihcanturk/pytorch_geometric.git; cd /opt/pyg/pytorch_geometric; pip install .;

I would also recommend testing in the nvidia pyg container for easiest setup
https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pyg/tags

[puririshi98]

based on testing with NVIDIA stack i think this is safe to merge

[puririshi98]

please fix linting and hopefuly other CI clears up(unrelated to you)