how transform adj of shape BN×BN to B×N×N

[zhinanchezhoubo]

I have a batch dataset, which is obtained from torch_geometric.loader, the dataset is such as: Batch(edge_attr=[2560, 4], edge_index=[2, 2560], x=[1154, 7], y=[64], batch=[1154], ptr=[65]), i want wo get the adj of the batch dataset, if i use torch_geometric.utils.to_dense_adj(edge_index,batch) will obtain a adj of 1154×1154，but i want to get the adj shape is B×N×N，B is subgraph number ,N is the max nodes number of subgraph. How can i do it?

[rusty1s]

Do you have a small example to reproduce? This works just fine for me:

from torch_geometric.data import Batch
from torch_geometric.utils import to_dense_adj
from torch_geometric.datasets import TUDataset

dataset = TUDataset('/tmp/TU', name='MUTAG')
batch = Batch.from_data_list(dataset[:10])
adj = to_dense_adj(batch.edge_index, batch.batch)
print(adj.shape)

[zhinanchezhoubo]

thank you much, great!

[zhinanchezhoubo]

How can i obtain the batch adj? for example:

from torch_geometric.data import Batch
from torch_geometric.utils import to_dense_adj
from torch_geometric.datasets import TUDataset

dataset = TUDataset('/tmp/TU', name='MUTAG')
batch = Batch.from_data_list(dataset[:10])
i want to obtain the btach's adj, whose seize is N×N，the N is the total number of nodes in batch. If i use 
adj = to_dense_adj(batch.edge_index, batch.batch)
print(adj.shape) #torch.Size([10, 28, 28])

But i want to the size of adj 166×166,166 is the batch's node number

[rusty1s]

Simply drop the batch argument in to_dense_adj in that case.

[zhinanchezhoubo]

import torch
from torch_geometric.datasets import TUDataset
from torch_geometric.data import Batch
from torch_geometric.utils import to_dense_adj

dataset = TUDataset(root='data/TUDataset', name='MUTAG')
graph = Batch.from_data_list(dataset)
print(graph)#DataBatch(edge_index=[2, 7442], x=[3371, 7], edge_attr=[7442, 4], y=[188], batch=[3371], ptr=[189])
adj = to_dense_adj(graph.edge_index)
print(adj.shape)#torch.Size([1, 3371, 3371])
adj = adj.squeeze()
print(adj.shape)#torch.Size([3371, 3371])

All right, thank you very much！

[rusty1s]

Super :)