Converting a scipy coo_matrix to pytorch sparse tensor

Question

I have a coo_matrix:

from scipy.sparse import coo_matrix
coo = coo_matrix((3, 4), dtype = "int8")

That I want converted to a pytorch sparse tensor. According to the documentation https://pytorch.org/docs/master/sparse.html it should follow the coo format, but I cannot find a simple way to do the conversion. Any help would be greatly appreciated!

Misuse

Answer 1

Using the data as in the Pytorch docs, it can be done simply using the attributes of the Numpy coo_matrix:

import torch
import numpy as np
from scipy.sparse import coo_matrix

coo = coo_matrix(([3,4,5], ([0,1,1], [2,0,2])), shape=(2,3))

values = coo.data
indices = np.vstack((coo.row, coo.col))

i = torch.LongTensor(indices)
v = torch.FloatTensor(values)
shape = coo.shape

torch.sparse.FloatTensor(i, v, torch.Size(shape)).to_dense()

Output

0 0 3
4 0 5
[torch.FloatTensor of size 2x3]

Broeker

Upon further investigation, I get a completely different result when I run your code: imgur.com/a/TwZHVhE I have no idea what is going on... – Misuse

Hmm— that seems very weird. Looks like the indices are getting mangled somehow. I can't reproduce it; I get the stated output. Maybe print debug the various components in a fresh notebook? – Broeker

It does, something goes awry with the indices, some float conversion I guess: imgur.com/a/QSPasIC – Misuse

My results (imgur.com/a/zJqqJs8) are still good with the latest package versions. Please let me know if you manage to debug the weirdness with the indices you're seeing on your side. – Broeker

i think it would be great if there was a pull request to do the conversion automatically from a coo matrix – Abseil

the indices = np.vstack(..) without dim assignment could be problematic, as coo.row could be multi-dimensional, thus giving rise to the issues commented above. – Cachucha

Answer 2

In case you want to convert a scipy.sparse.csr_matrix to a torch.sparse_coo_tensor, you can do it the following way:

import torch
from scipy.sparse import csr_matrix

csr = csr_matrix([[1, 2, 0], [0, 0, 3], [4, 0, 5]])

# Convert to PyTorch sparse tensor
pt_tensor = torch.sparse_coo_tensor(csr.nonzero(), csr.data, csr.shape)

Output:

tensor(indices=tensor([[0, 0, 1, 2, 2],
                       [0, 1, 2, 0, 2]]),
       values=tensor([1, 2, 3, 4, 5]),
       size=(3, 3), nnz=5, layout=torch.sparse_coo)

Extend

Answer 3

import torch
import numpy as np
from scipy.sparse import coo_matrix

coo = coo_matrix((3, 4), dtype = "int8")
row = torch.from_numpy(coo.row.astype(np.int64)).to(torch.long)
col = torch.from_numpy(coo.col.astype(np.int64)).to(torch.long)
edge_index = torch.stack([row, col], dim=0)

#Presuming values are floats, can use np.int64 for dtype=int8
val = torch.from_numpy(coo.data.astype(np.float64)).to(torch.float)

out = torch.sparse.FloatTensor(edge_index, val, torch.Size(coo.shape)).to_dense() 

Cachucha