awslabs / awslabs/graphstorm

Questions on using GS Python API for inference with new graphs and subgraphs

Open
#1,339 13 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
452
Forks
76
PR merge metrics
No merged PRs in 30d

Description

Hi, I have trained a link prediction model on a large original graph using the GS CLI (`graphstorm.gconstruct.construct_graph` & `graphstorm.run.gs_link_prediction`).

Now I am trying to switch to the GS Python API to perform the following steps:

1. Load the trained model.

2. Prepare new graph data (more than 1M small graphs).

3. Prepare subgraph data (more than 1M subgraphs, each consisting of a subgraph from the original graph plus some new nodes and edges).

4. Compute node embeddings for both new graph data and subgraph data.

Here are my questions for each step:

Q1. Model loading & feature dimensions
I load the trained model with the following code. However, the model only contains `gnn` and `embed` parameters, which take 128-dim features as input. Since my original node/edge data only has 5-dim features, how should I transform my new data so that it matches the model input?
```
model = gs.create_builtin_lp_gnn_model(train_data.g, config, train_task=False)
model.restore_dense_model('model_path/epoch-2')
```

Q2. Graph construction performance
I need to process a large number of new graphs. Using graphstorm.gconstruct.construct_graph is extremely slow (processing a graph with 5 nodes takes roughly the same time as a graph with 1M nodes: 20–30 minutes).
Is there a way to accelerate this process?
Alternatively, can I directly use a DGLGraph (instead of DistGraph or GSgnnData) for inference on small graphs?

Q3. Handling subgraphs with new nodes/edges
I need to run inference on subgraphs that consist partly of the original graph and partly of new nodes/edges (each node has a unique ID).
How should I represent this data so that the model can distinguish between the “original” subgraph part and the “new” part?

Q4. Extracting node embeddings
Finally, what’s the recommended way to obtain node embeddings for subgraphs (i.e., a subgraph plus newly added nodes and edges)?

It would be very helpful if you could also provide some sample code for these steps.

Thanks in advance for your help!

Contributor guide

Open the contributing guide

Research direction

Start by reading graphstorm.gconstruct.construct_graph and graphstorm.run.gs_link_prediction, then inspect gs.create_builtin_lp_gnn_model, restore_dense_model, DistGraph, GSgnnData, and DGLGraph inference paths. Document the supported model-loading, graph and subgraph representation, performance, and node-embedding workflows, with sample code that answers Q1–Q4.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
data-engineering, machine-learning
Issue type
Documentation
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
20/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.