batch of graphs for GAT - ensuring the same node in one graph is the same for another graph

[amjass12]

Hi all,

I am training a GAT on a batch of 10 individual graphs and classifying each one in to their respective communities as determined by Louvain-clustering. The node names are strings and some graphs contain more nodes than others. One thing that is important though is that there are similar nodes in different graphs. For example 'person A' is present in graph one, graph two, graph three etc... so the mapping (and thus embedding representation) clearly has to be common. But i realise this mapping might not be preserved in the dataset preparation script. Any help is appreciated. My dataset prep is as follows (I am using w2v vectors trained on these graphs individually as node features for each node in each graph):

graph_dataset_list = []
for i, j, k in zip(files, w2v_reordered, community_reordered):
    splitter = T.Compose([T.RandomNodeSplit()])

    graph = nx.read_gpickle(os.path.join('data/gpickles/', i))
    model = Word2Vec.load(os.path.join('data/models/', j))
    comm = pd.read_csv(os.path.join('data/comms/', k), index_col=0)

    df = pd.DataFrame(list(graph.nodes), columns=['genes'])
    orders = np.array_equal(comm.index.values, df['genes'])

    word_vecs = []
    for word in list(graph.nodes):
        vec = model.wv[word]
        word_vecs.append(vec)

    final_word_vecs = np.array(word_vecs)

    node_degree = [val for (node, val) in graph.degree()]
    node_degree = np.round(np.array(node_degree) / len(graph.nodes), 3)

    final_node_features = np.c_[final_word_vecs, node_degree]
    print(final_node_features.shape)

    final_node_features = torch.Tensor(final_node_features)

    G = nx.convert_node_labels_to_integers(graph)
    G = G.to_directed()
    edges = list(G.edges)
    edge_index = torch.tensor(edges, dtype=torch.long).t().contiguous()

    y = comm['community'].to_numpy()
    print(f'shape of y labels: {y.shape}')
    y = torch.from_numpy(y)

    torch_graph = Data(x=final_node_features, y=y, edge_index=edge_index)
    torch_graph = splitter(torch_graph)

    graph_dataset_list.append(torch_graph)

The reason this is important is because, 'person A' might be node 0 in one graph, but node 19 in another graph, and so when converting from integers in the loop of each graph, this mapping is not preserved

I would appreciate any advice or solution to this. Many thanks!!

[rusty1s]

I think the most reasonable thing to do would be to maintain a global_node_idx for every node in your graph, that maps similar nodes to the same node id. You can then use this id to query an Embedding layer, e.g., embedding(data.global_node_idx).

[rusty1s]

I don't think you need to re-index the graph. The graph can be arbitrarily ordered. If you have node 0 in graph 0, and node 3 in graph 1 referencing the same node, then data0.global_node_idx[0] and data1.global_node_idx[3] should point to the same index.

[amjass12]

Hi @rusty1s

thank you very much - Can you explain why this doesn't mess up the embedding layers though of the GAT? shouldn't the embedding of node 0 and node 3 in your example be the same?

In my mind, this is the same as in the canonical example of transformers for language - the embedding for word 'hello' would be the same in two different training examples right? wherever they appear in the sentence? - the updates for the embedding of hello happen on the same vector? this is one thing that is confusing me with the graphs

[rusty1s]

Yes, this is exactly what would happen here, right? Maybe I wasn't clear. Let's assume the word "hello" is encoded in the first graph at position 2, and the same word is encoded in the second graph at position 4, and it reference position 6 in the embedding matrix. Then you would have

data1.global_node_idx = [dummy, 6, ...]
data2.global_node_idx = [dummy, dummy, dummy, 6, ...]

and in your forward pass, you would do

embedding = torch.nn.Embedding(...)

def forward(self, global_node_idx):
    x = embedding(global_node_idx)

[amjass12]

Thank you @rusty1s

I am sorry for the confusion - I think i am confused because - I convert each networkx graph to integers as per this line of code:
G = nx.convert_node_labels_to_integers(graph)

So the mapping if effectively lost and this is where i became really confused. Is there a way to go around this and maintain the strings so that it can now achieve the behaviour you describe?

The reason this matter particularly for me is because i need to retrieve the embedding of an individual node after the GAT has been trained - which i have been achieving like:

whole_graph_emb = model.gat_conv1(node_features, edges_index)
embedding_node_x = whole_graph_emb[node_idx]

i realise now however, that for example, i am not sure what index 0 would refer to? index 1 etc... i am not sure which nodes they correspond to:

Is there a page that describes how to take care of batches in this way? thank you!

[rusty1s]

I would assume you can just add an additional attribute to your networkx graph that maintains the mapping (which should be preserved by from_networkx). Let me know if this helps.

[amjass12]

Hi @rusty1s ,

ok got it so to clarify - as long as 'node_red' is mapped to for example, the integer 0 in graph one and the same in graph two - then this node should have an identical mapping between both graphs when 'from networkx is run'right? so that now, the embedding for this node is the same in the global neural network?

Then i assume the standard approach:

class GATmodel(torch.nn.Module):
    def __init__(self, input_dim, hidden_channels, num_heads):
        super().__init__()
        self.gat_conv1 = GATConv(in_channels=input_dim, out_channels=hidden_channels, heads=num_heads)
        self.gat_conv2 = GATConv(in_channels=hidden_channels*num_heads, out_channels=53, heads=1)

    def forward(self, data=None, node_features=None, edges_index=None):
        if not data:
            x, edge_index = node_features, edges_index
        else:
            x, edge_index = data.x, data.edge_index
        
        x = F.dropout(x, p=0.4, training=self.training)
        x = self.gat_conv1(x, edge_index)
        x = F.elu(x)
        x = F.dropout(x, p=0.4, training=self.training)
        x = self.gat_conv2(x, edge_index)
        return x

can be applied? many thanks!

[rusty1s]

Yes :)

[amjass12]

Thank you so much @rusty1s !

I wonder if it might be possible to preprocess and concatenate batches of graphs and then have a function that can effectively 'tokenize' nodes? and therefore create a mapping across all batches of graohs - i dont know if this would work, but would but would be a convenient function!

[rusty1s]

Do you have an example of what you mean? :)

[amjass12]

Hi @rusty1s

apologies for the delay in replying - work has been very busy. I meant more the canonical way in which a corpus of text is tokenized in traditional NLP.

So for example, Once all graphs are concatenated in a batch - then might be a good time to create a mapping of nodes to integers so that all nodes that share the same name are given the same index, so that their embedding vector is at the same position in the GAT - this way the confusion of shared nodes would be taken away and the user would not have to worry about re-arranging the names in different graphs so that common nodes share the same index.

[rusty1s]

I think it might be easier if you create that shared mapping beforehand for all nodes in all your graphs. In general, batching just performs a concatenation of the "word indices", and computing the mapping for every batch sounds expensive. Let me know if I am misunderstanding.