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

[han-yue-1523]

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!

[classicsong]

I think you are looking for something similar to https://graphstorm.readthedocs.io/en/latest/cli/model-training-inference/real-time-inference.html.

Maybe you can find something useful from it.

[han-yue-1523]

I think you are looking for something similar to https://graphstorm.readthedocs.io/en/latest/cli/model-training-inference/real-time-inference.html.

Maybe you can find something useful from it.

I am not using SageMaker. Can I do it on my own machine?

[zhjwy9343]

If not using SageMaker, you will need to maintain a 7*24 running server that hosts the trained model for real-time inference. SageMaker endpoint uses TorchServ, which you can borrow for using it on your own machine.

For the model loading and prediction, you can refer the https://github.com/awslabs/graphstorm/blob/main/sagemaker/run/realtime_entry_points/node_prediction_entry.py code.

[han-yue-1523]

I don't need real-time inference because all the data is saved on my local machine. I am confused about feature dimensions, data type, and data processing time.

[zhjwy9343]

May I ask why you want to do this rather than prepare a giant graph that includes all new nodes?

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

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

[zhjwy9343]

In general, if you prepare a new graph data (>1M small graphs), you can do LP inference directly on the new graph. No need to handle the subgraphs.

For the inference CLI, here is an example: https://catalog.workshops.aws/graphstormkdd2023/en-US/standalone-labs/standalone-lab3#3.-launch-link-prediction-training-and-inference

[han-yue-1523]

My task involves new nodes that are associated differently across subgraphs, and these nodes are only visible within their respective subgraphs. This requires me to handle both small graphs and subgraphs. Is GraphStorm suitable for this scenario?

[zhjwy9343]

My understanding is that you can treat these subgraphs as disconnected components of a giant graph. So you can just prepare all data into one giant graph by using construct CLI, and then run inference on the giant graph once.

Same thing, you can train your model in the same way.

[han-yue-1523]

That’s a great idea! However, in my case, there are two other challenges:

(1) Different subgraphs may share the same nodes, so if I create a large disconnected graph, I would need to duplicate many nodes, which could significantly increase the graph size.

(2) I want to leverage neighbor information from the original graph when working with subgraphs, rather than treating subgraphs as fully independent.

Do you have any suggestions? Thank you.

[zhjwy9343]

(1) Different subgraphs may share the same nodes, so if I create a large disconnected graph, I would need to duplicate many nodes, which could significantly increase the graph size.
If they are the same nodes in the physical world, why put them indifferent subgraphs? Should they be the same nodes that can help to build connections? In terms of graph size, as long as you have the computational resources, GraphStorm can handle +100M of nodes easily.

(2) I want to leverage neighbor information from the original graph when working with subgraphs, rather than treating subgraphs as fully independent.
In such case, you just put these subgraphs back into the original graph, and let them connected.

I don't see any need of subgraphs based on our conversation. What you have is a typical training graph versus inference graph case. That is, we build a graph for training, and then add some new nodes/edges to the training graph to form an inference graph. Then we use the trained model to do inference on the inference graph. GraphStorm can fully support this case.

[han-yue-1523]

Let me clarify my task with an example below. Please note that in practice, the original graph G has more than 10M nodes, and the number of inference cases k is also greater than 10M.

Original Graph G: 1000 nodes (N1 … N1000) + edges.

Inference cases:

(1) G + N1001 + new edges (only need embeddings for N1, N2, N1001)

(2) G + N1002 + N1003 + new edges (only need embeddings for N2, N5, N1002, N1003)
...
(k) G + N2000 … N3000 + new edges (only need embeddings for N100, N200, N333, N999, N2000 … N3000)
...

Key requirement:
For each inference case, message passing must be limited to the original graph plus the new nodes/edges of that case only.

For example: In case (1), the graph should be G + N1001 without seeing N1002, N1003,…, and in case (2), the graph should be G + N1002 + N1003 without seeing N1001, N1004, N1005, ….

Given the key requirement above, I cannot simply add all new nodes into G. Since k is very large and each inference only requires embeddings for a small portion of G, I am looking for a more efficient way to construct graphs and obtain embeddings than the standard inference approach. (See Q2)

[zhjwy9343]

I can see two solutions for your case:

1. A straightforward solution: extract the small subgraph for each inference case, and then build a giant graph, which might include ~100M nodes. GraphStorm can handle this scale. But in this solution, the same node could have different embeddings, e.g., the N2 node that appears in both inference case 1 and 2, because the subgraphs are different.
2. A custom solution: create a custom sampler like this one (https://github.com/awslabs/graphstorm/blob/main/python/graphstorm/dataloading/sampler.py#L479). You can set some edge features to limit the edges to be sampled for each inference case. For example, set the N001 to have a time and all the edges linked to N001 to have the same time. And then in the sampler only sample edges with this time and those without time, which is in the G.

[han-yue-1523]

Thank you, I’ll give it a try. Regarding the first solution, I noticed that the provided examples here train and run inference on the same dataset (using different data loaders). How can I instead run inference with a previously trained model on a new dataset?

Please see the code below.

train_dataset = GSgnnData(train_part_config)
train_embeddings = do_mini_batch_inference(model, train_dataset, batch_size=1024, fanout=[10, 10])

test_dataset = GSgnnData(test_part_config)
test_embeddings = do_mini_batch_inference(model, test_dataset, batch_size=1024, fanout=[10, 10])

The code runs correctly on train_dataset but fails on test_dataset. Both datasets have the same node and edge types as well as identical feature lists, but they contain different nodes and edges (and test_dataset may potentially not covering all node/edge types in train_dataset).