Different adjacency matrix per sample

[javirk]

Using GCNConv, I am trying to feed each sample in the batch with a different adjacency matrix. If all of them had the same, I would simply do:

class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = GCNConv(10, 16)

    def forward(self, x, edge_index, edge_weight):
        return self.conv1(x, edge_index, edge_weight)

A = torch.randint(low=0, high=2, size=(6, 6))
edge_index_sparse = A.to_sparse()

X = torch.rand((4, 6, 10))
model = Net()
out = model(X, edge_index_sparse.indices(), edge_index_sparse.values())

For a graph with 6 nodes, a batch size of 4, and where the input tensors have 10 features.

However, if I try to initialize a different adjacency per sample:

A = torch.randint(low=0, high=2, size=(4, 6, 6))

I get an error saying that the edge_index matrix has three components in the first dimension (of course):

assert edge_index.size(0) == 2
AssertionError

Is there any way to do this? Even with a different type of layer, that's not an issue.

[rusty1s]

PyG does not support traditional batching via a dedicated batch dimension in order to easily support mini-batching of examples of varying size. As such, edge_index needs to be a [2, num_edges] matrix even in mini-batch scenarios by stacking adjacency matrices diagonally (see the documentation for more information). When having a list of data objects holding x and edge_index tensors, you can simply do:

batch = Batch.from_data_list(data_list)
conv(batch.x, batch.edge_index)

[javirk]

Thanks for the reply. In my case, I don't have a data object but pytorch tensors because they fit better into my pipeline. I implemented something similar to what you propose by stacking the adjacency matrix diagonally and making a block matrix. Given a batch of adjacency matrix adj_mat with dimensions [B, N, N], you can transform it into a block matrix with:
adj_mat = torch.block_diag(*adj_mat)

Then the to_sparse() method works just fine and the edge_index matrix is [2, num_edges]:
adj_sparse = adj_mat.to_sparse()
This can be fed to the convolution.

Thanks!

[rusty1s]

Super, I didn't know about torch.block_diag :) Keep in mind that this will consume a lot of memory for dense matrices and should only be used for testing purposes.