pytorch 邻接矩阵转稀疏矩阵 (dense matrix to coo matrix)

本文已参与「新人创作礼」活动，一起开启掘金创作之路。

如何将一个dense矩阵，比如：一个 $N \times N$

import numpy as np
import torch
a = np.array([[0, 1.2, 0],[2, 3.1, 0],[0.5, 0, 0]])
idx = a.nonzero() # (row, col)
data = a[idx]
# to torch tensor
idx_t = torch.LongTensor(np.vstack(idx))
data_t = torch.FloatTensor(data)
coo_a = torch.sparse_coo_tensor(idx_t, data_t, a.shape)
print(coo_a)
import scipy.sparse as sp
import numpy as np
import torch
a = np.array([[0, 1.2, 0],[2, 3.1, 0],[0.5, 0, 0]])
coo_np = sp.coo_matrix(a)
data = coo_np.data
idx_t = torch.LongTensor(np.vstack((coo_np.row, coo_np.col)))
data_t = torch.FloatTensor(data)
coo_a=torch.sparse_coo_tensor(idx_t,data_t,a.shape)
方法三（从tensor到sparse tensor）：
import numpy as np
import torch
a = torch.tensor([[0, 1.2, 0],[2, 3.1, 0],[0.5, 0, 0]])
idx = torch.nonzero(a).T  # 这里需要转置一下
data = a[idx[0],idx[1]]
coo_a = torch.sparse_coo_tensor(idx, data, a.shape)
print(coo_a)
如果想转回dense的tensor
那么可以直接调用自带的方法：
coo_a.to_dense()
什么是COO，CSR矩阵
COO矩阵简介：
indptr指向索引和数据中的行开始
长度为n_row + 1，最后一项 = 值数 = 索引和数据的长度
第i行的非零值是data[indptr[i]:indptr[i+1]] 和列索引indices[indptr[i]:indptr[i+1]]
item (i, j)可以作为data[indptr[i]+k]访问，其中k是j在indices[indptr[i]:indptr[i+1]]中的位置