How to pass batch of Graphs in GATv2Conv?

Question

I have a batched input to GATv2Conv with node matrix of shape [batch_sz , num_nodes , node_feature_dim] , but the GATv2Conv accepts input of dim 2 ,searching through the internet , I found some solution ... (not the one I want)

data_list = [Data(x= torch.squeeze(torch.index_select(x, dim= 0, index = torch.tensor([idx]))) , 
                              edge_index= self.edge_indices ,
                              edge_attr= torch.squeeze(torch.index_select(adj_mats[i], dim= 0, index = torch.tensor([idx])))) 
                         for idx in range(self.batch_sz)] 
            batch  = Batch.from_data_list(data_list)

But using above solution , the distinction between graphs got lost , becuase : batch.x.shape gave [batch_sz * num_nodes , node_feature_dim]...

It simply put all nodes of all graphs in one single graphs.. Now there is shared calculations between different graphs , which is strictly undesirable... As when applying some graph pooling layer, I don't know which nodes belonged to which graph....

Pls suggest some fix for this issue ...

Thanks in advance

Answer 1

For efficiency reasons, it is the standard for graph learning libraries to use the format [n_batches x num_nodes, node_feature_dim], instead of having batches in an extra batch dimension. Specifically, the different graphs of a batch are transformed into a single large graph with disjoint subgraphs. However, because there is no connection between any two graphs of the batch, neighborhood operations like message-passing will not leak information into different graphs of the batch. This is not true for global graph operations, such as pooling. Such operations require a batch-index as an additional input. The batch-index maps each node onto its batch index. In the following, $f$ is a global graph operation (set operation), $N$ is note number of graphs $\mathcal{G}$ in your batch.

$$ X = \{ f\left(\mathcal{G}_i\right) : \forall i \in [0,..,N] \} $$

Each result will only contain information from within the respective batch. In the Batch class that you are using, according to the documentation, there is a batch field that tells you which nodes belong to which graph:

In addition, single graphs can be identified via the assignment vector batch, which maps each node to its respective graph identifier.

For more on graph batching in general here is a colab notebook that contains some nice demo code.