scipy和pytorch的sparse库有啥区别

638 阅读1分钟

这个问题具体是针对coo matrix这个存储方式问的

sparse库

首先sparse库是用来存储稀疏矩阵的,pytorch和scipy都有这么个库(好像pytorch比较高的版本才有)。sparse库最主要的就是稀疏矩阵的表示,一般使用三元组的方式存储:(行,列,值),比如scipy的 coo_matrix((data, (i, j)), [shape=(M, N)])

pytorch和scipy的sparse有啥区别

因为看到GNN一些代码里在处理邻接矩阵的时候会把pytorch的sparse转换成scipy的,做个normalization再变回pytorch的sparse,不知道为啥要这么做。

原因:scipy的sparse支持很多sum, max这种操作,但是老版本的pytorch的还不支持。即二者支持的功能不一致。

关于具体支持的内容可以查看官方文档,拉到最底下就是:

scipy: scipy.sparse.coo_matrix pytorch: pytorch sparse

pytorch的进化速度特别快,1.10的时候已经支持稀疏矩阵相乘了torch.sparse.matmul,包括上面提到的sum, max啥的操作

相互转换

贴一下二者怎么相互转换的吧

def sparse_mx_to_torch_sparse_tensor(sparse_mx):
    """Convert a scipy sparse matrix to a torch sparse tensor."""
    sparse_mx = sparse_mx.tocoo().astype(np.float32)
    indices = torch.from_numpy(
        np.vstack((sparse_mx.row, sparse_mx.col)).astype(np.int64))
    values = torch.from_numpy(sparse_mx.data)
    shape = torch.Size(sparse_mx.shape)
    return torch.sparse.FloatTensor(indices, values, shape)


def torch_sparse_tensor_to_sparse_mx(torch_sparse):
    """Convert a torch sparse tensor to a scipy sparse matrix."""


    m_index = torch_sparse._indices().numpy()
    row = m_index[0]
    col = m_index[1]
    data = torch_sparse._values().numpy()

    sp_matrix = sp.coo_matrix((data, (row, col)), shape=(torch_sparse.size()[0], torch_sparse.size()[1]))

    return sp_matrix
    ```