import numpy as np
import scipy.sparse as sp
import torch
a = np.array([[1,0,0,1,0],[0,1,1,0,0],[0,0,0,0,1],[1,0,0,0,1],[0,0,1,0,0]])
a = sp.coo_matrix(a,shape=(5,5))
sparse_mx = a.tocoo().astype(np.float32)
print(sparse_mx.row)
print(sparse_mx.col)
print(sparse_mx.data)
print(sparse_mx.shape)
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)
print(torch.sparse.FloatTensor(indices, values, shape))
输出:
[0 0 1 1 2 3 3 4]
[0 3 1 2 4 0 4 2]
[1. 1. 1. 1. 1. 1. 1. 1.]
(5, 5)
tensor(indices=tensor([[0, 0, 1, 1, 2, 3, 3, 4],
[0, 3, 1, 2, 4, 0, 4, 2]]),
values=tensor([1., 1., 1., 1., 1., 1., 1., 1.]),
size=(5, 5), nnz=8, layout=torch.sparse_coo)