Easiest/efficient way to implement a unidirectional bipartite graph

[ssharpe42]

I am trying to model a very simple unidirectional bipartite graph where there are src nodes and dst nodes. I want to predict only on the dst nodes and there is a many to one (src to dst) relationship.

Here is an example edge list with 9 src nodes and 3 dst nodes.

src = torch.tensor([0, 1, 2,  3, 4, 5, 6, 7, 8])
dst =  torch.tensor([0, 0, 0,  1, 2, 2, 2, 2, 2])

In this case I would like to do a graph convolution (Sage, GAT, etc) on src nodes 0,1,2 and the dst node 0 to make a prediction for dst node 0 and so on...

Is there a way to do this easily with pytorch_geometric or would it be better to just implement from scratch?

This is how you can do it in dgl:

# Src and dst nodes & create graph
src = torch.tensor([0, 1, 2,  3, 4, 5, 6, 7, 8])
dst =  torch.tensor([0, 0, 0,  1, 2, 2, 2, 2, 2])
g = dgl.heterograph({('src_type', 'edge', 'dst_type') : (src, dst)})

# Generate node features and labels (dst_nodes have different number of features)
g.srcdata["node_feat"] = torch.rand(len(g.srcnodes()), 10)
g.dstdata["node_feat"] = torch.rand(len(g.dstnodes()), 3)


conv = dgl.nn.SAGEConv((10 , 3), 8, aggregator_type = "mean")
fc = nn.Linear(8, 1)
h= conv(g, (g.srcdata["node_feat"], g.dstdata["node_feat"]))
print(h)

tensor([[ 0.4699, -0.0037,  0.0341,  1.3000,  0.5906,  0.9941, -0.2148, -0.8874],
        [-0.0616,  0.4049,  0.0663,  0.7090,  0.4139,  1.4098, -0.3871, -0.3437],
        [ 0.2171, -0.6719,  0.5209,  1.3910,  0.9904,  1.0238,  0.6527, -0.3474]],
       grad_fn=<AddBackward0>)

[rusty1s]

Yes, this actually is implemented quite similar to DGL:

from torch_geometric.nn import SAGEConv

src = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7, 8])
dst =  torch.tensor([0, 0, 0, 1, 2, 2, 2, 2, 2])
edge_index = torch.stack([src, dst], dim=0)

x_src = torch.rand(9, 10)
x_dst = torch.rand(3, 3)

conv = SAGEConv((10, 3), 8)
out_dst = conv((x_src, x_dst), edge_index)

Note that in PyG master, you can now store this bipartite graph via a HeteroData object as well, which allows you to do sampling/mini-batching and to apply transformations:

from torch_geometric.data import HeteroData

data = HeteroData()
data['src'].x = x_src
data['dst'].x = x_dst
data['src', 'dst'].edge_index = edge_index

from torch_geometric.transforms import ToUndirected

data = ToUndirected()(data)

[ssharpe42]

Great thanks! Maybe I was overthinking a bit. Looks like HeteroData makes things pretty easy for batching.

If I was to just use the first approach and stay on my installed release ... just retrieving (src, dst,edge_index) indices aligning with specific source nodes would do the trick since it is such a simple problem right?

[rusty1s]

Yes, working on bipartite graphs is a little bit simpler as you do need need to take into account the aggregation of features from different edge types pointing to the same destination.

Nonetheless, keep in mind that you need to perform message passing in both directions (by transposing edge_index), as other-wise only destination nodes will update their representations.

[ssharpe42]

I actually don't care about dst-->src messages. I only want unidirectional so the output dimensions of the conv should be [dst sample size, hidden size].

All src nodes are unique and so are all dst nodes so they basically each form their own 1 hop neighborhood graph where the src flow into the dst.

Maybe as you suggest HeteroData is the best solution for this. I just didn't know if there was a more elegant solution than this:

from torch_geometric.utils import k_hop_subgraph, subgraph
from torch_geometric.nn import SAGEConv

src = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7, 8])
dst =  torch.tensor([0, 0, 0, 1, 2, 2, 2, 2, 2])
edge_index = torch.stack([src, dst], dim=0)

x_src = torch.rand(9, 10)
x_dst = torch.rand(3, 3)

conv = SAGEConv((10, 3), 8)

# sample of dst nodes
dst_batch = torch.tensor([1,2])
_, sub_edge_index, _, _ = k_hop_subgraph(dst_batch, 1,  edge_index, relabel_nodes=False, num_nodes=None, flow='source_to_target')

# map nodes to start at 0
sub_src = sub_edge_index[0]
mapping_src = torch.empty(sub_src.max()+1, dtype = torch.long) 
mapping_dst = torch.empty(dst_batch.max()+1,dtype = torch.long)

mapping_src[sub_src] = torch.arange(sub_src.size(0))
mapping_dst[dst_batch]= torch.arange(dst_batch.size(0))

sub_edge_index[0] = mapping_src[sub_edge_index[0]]
sub_edge_index[1] = mapping_dst[sub_edge_index[1]]

# apply conv to subset of dst nodes and their src
conv((x_src[sub_src], x_dst[dst_batch]), sub_edge_index )

tensor([[-0.3151, -0.3093,  0.1171,  0.1739,  0.9588,  0.0307,  0.0856,  0.2794],
        [-0.0697, -0.4284,  0.2630,  0.1877,  0.8047,  0.3782,  0.0553,  0.4039]],
       grad_fn=<AddBackward0>)

[rusty1s]

Your implementation is correct, but keep in mind that without reverse edges you are limited to 1-layer GNNs/learning from the direct neighborhood of source nodes. In practice, it might be also beneficial to learn from larger receptive fields, e.g., from destination nodes that share similar source nodes. In order to allow that, you need to extend your data to account for reverse edges as well. In its most simple form, this looks similar to the following:

x_src = ...
x_dst = ...
src_to_dst_edge_index = ...
dst_to_src_edge_index = src_to_dst_edge_index.flip(0)

conv1_src_to_dst = SAGEConv((10, 3), 16)
conv1_dst_to_src = SAGEConv((3, 10), 16)

conv2_src_to_dst = SAGEConv((16, 16), 8)

x_dst = conv1_src_to_dst((x_src, x_dst), src_to_dst_edge_index).relu()
x_src = conv1_dst_to_src((x_dst, x_src), dst_to_src_edge_index).relu()
x_dst = conv2_src_to_dst((x_src, x_dst), src_to_dst_edge_index)