Qubix
Qubix

Reputation: 4353

Pytorch Geometric sparse adjacency matrix to edge index tensor

My data object has the data.adj_t parameter, giving me the sparse adjacency matrix. How can I get the edge_index tensor of size [2, num_edges] from this?

Upvotes: 5

Views: 8492

Answers (2)

ndrwnaguib
ndrwnaguib

Reputation: 6135

You can use torch_geometric.utils.convert.from_scipy_sparse_matrix.

>>> from torch_geometric.utils.convert import from_scipy_sparse_matrix
>>> edge_index = torch.tensor([
...    [0, 1, 1, 2, 2, 3],
...    [1, 0, 2, 1, 3, 2],
>>> ])
>>> adj = to_scipy_sparse_matrix(edge_index)
>>> # `edge_index` and `edge_weight` are both returned
>>> from_scipy_sparse_matrix(adj)
(tensor([[0, 1, 1, 2, 2, 3],
        [1, 0, 2, 1, 3, 2]]),
tensor([1., 1., 1., 1., 1., 1.]))

Upvotes: 1

Berriel
Berriel

Reputation: 13641

As you can see in the docs:

Since this feature is still experimental, some operations, e.g., graph pooling methods, may still require you to input the edge_index format. You can convert adj_t back to (edge_index, edge_attr) via:

row, col, edge_attr = adj_t.t().coo()
edge_index = torch.stack([row, col], dim=0)

Upvotes: 1

Related Questions