Using one node's features as the READOUT function for graph level classification problem

Questions
#1

Hi,

Thanks for this great open-sourced package.

I am trying to do graph level classification following https://docs.dgl.ai/en/0.4.x/tutorials/basics/4_batch.html.

It uses dgl.mean_nodes() to represent each graph in a batch and fed into the downstream classifier. However, in my specific use case, I am interested in representing this graph with one node’s feature in the graph, I know the index of the node and I could retrieve it by using h[index] for one graph but I am wondering how do I do it in the batched graph setting?

Thank you so much!

#2

You only need to slice the rows of the node features in the batched graph with the node indices.

#3

Thanks for the prompt reply. How do I find the way DGL separates each graph in the batched graph? because each graph is associated with a different node index that I can slice.

#4

See the example below:

import dgl
import torch
from dgl import DGLGraph

g1 = DGLGraph()
g1.add_nodes(2)
g1.ndata['h'] = torch.tensor([[0.], [1.]])
g2 = DGLGraph()
g2.add_nodes(3)
g2.ndata['h'] = torch.tensor([[2.], [3.], [4.]])
bg = dgl.batch([g1, g2])
print(bg.batch_num_nodes)
# [2, 3]

# Node indices to fetch, node 0 from the first graph 
# and node 1 from the second graph
to_fetch = torch.LongTensor([0., 1.])
num_nodes_ = bg.batch_num_nodes
num_nodes_.insert(0, 0)
offset = torch.cumsum(torch.LongTensor(num_nodes_), dim=0)[:-1]
print(bg.ndata['h'][to_fetch + offset])
# tensor([[0.], [3.]])
#5

Works well. Thank you!!