【稀疏矩阵】使用torch.sparse模块

文章目录

稀疏矩阵的格式
- coo
- csr
- csc
[Construction of Sparse COO tensors](#Construction of Sparse COO tensors)
[Construction of CSR tensors](#Construction of CSR tensors)
[Linear Algebra operations（稀疏与稠密之间混合运算）](#Linear Algebra operations（稀疏与稠密之间混合运算）)
[Tensor methods and sparse（与稀疏有关的tensor成员函数）](#Tensor methods and sparse（与稀疏有关的tensor成员函数）)
- coo张量可用的tensor成员函数（经实测，csr也有一些可以用，比如dim()）
[Torch functions specific to sparse Tensors（与稀疏有关的torch函数）](#Torch functions specific to sparse Tensors（与稀疏有关的torch函数）)
支持稀疏张量的常规torch函数
支持稀疏张量的一元函数

稀疏矩阵的格式

目前，torch.sparse和scipy.sparse模块比较支持的主流的稀疏矩阵格式有coo格式、csr格式和csc格式，这三种格式中可供使用的API也最多。

coo

将矩阵中非零元素的坐标和值分开存储在3个数组中，3个数组长度必须相同，表示有n个非零元素。

csr

分 Index Pointers、Indices、Data3个数组存储。

Index Pointers：第 i个元素记录这个矩阵的第 i行的第1个非零值在 Data数组的起始位置，第 i+1个元素记录这个矩阵的第 i行的最后一个非零值在 Data数组的终止位置（不包含右边界）。因此，这个矩阵的行数等于 len(Index Pointers)-1，第 i行非零值的个数等于 Index Pointers[i+1]-Index Pointers[i]。
Indices：第 i个元素记录这个矩阵的第 i个非零值的列坐标。
Data：第 i个元素记录这个矩阵的第 i个非零值的具体数值，排列顺序严格按照行优先，列次先。

csc

与csr唯一的不同在于列优先，其他规则一模一样。

Construction of Sparse COO tensors

常规构建

python 复制代码

>>> i = [[0, 1, 1],
         [2, 0, 2]]
>>> v =  [3, 4, 5]
>>> s = torch.sparse_coo_tensor(i, v, (2, 3))
>>> s
tensor(indices=tensor([[0, 1, 1],
                       [2, 0, 2]]),
       values=tensor([3, 4, 5]),
       size=(2, 3), nnz=3, layout=torch.sparse_coo)
>>> s.to_dense()
tensor([[0, 0, 3],
        [4, 0, 5]])

torch中，稀疏矩阵的存储方式记录在 tensor.layout中，可以通过检查 torch.layout == torch.sparse_coo来判断是否是coo张量。此外，稠密张量的 layout等于 strided。

稠密混合的coo张量

python 复制代码

>>> i = [[0, 1, 1],
         [2, 0, 2]]
>>> v =  [[3, 4], [5, 6], [7, 8]]
>>> s = torch.sparse_coo_tensor(i, v, (2, 3, 2))
>>> s
tensor(indices=tensor([[0, 1, 1],
                       [2, 0, 2]]),
       values=tensor([[3, 4],
                      [5, 6],
                      [7, 8]]),
       size=(2, 3, 2), nnz=3, layout=torch.sparse_coo)

此方案与常规的coo构建方式不同，values中每个元素可以是一个向量，表示对应坐标的稠密张量，因此，创建出的coo张量也多出了一个维度。

带有重复坐标的coo张量

python 复制代码

>>> i = [[1, 1]]
>>> v =  [3, 4]
>>> s=torch.sparse_coo_tensor(i, v, (3,))
>>> s
tensor(indices=tensor([[1, 1]]),
       values=tensor(  [3, 4]),
       size=(3,), nnz=2, layout=torch.sparse_coo)
>>> s.to_dense()
tensor([0, 7, 0])

如果输入的坐标有重复，则创建出的coo张量会自动把坐标重复的元素值相加。此外，可以通过成员函数 .coalesce()把重复坐标的元素值相加，将这个coo转换成一个不重复的张量；也可以通过 .is_coalesced()检查这个coo是否存在重复的坐标。

Construction of CSR tensors

按照 Index Pointers、Indices、Data三个数组的定义构建即可。

python 复制代码

>>> crow_indices = torch.tensor([0, 2, 4])
>>> col_indices = torch.tensor([0, 1, 0, 1])
>>> values = torch.tensor([1, 2, 3, 4])
>>> csr = torch.sparse_csr_tensor(crow_indices, col_indices, values, dtype=torch.float64)
>>> csr
tensor(crow_indices=tensor([0, 2, 4]),
       col_indices=tensor([0, 1, 0, 1]),
       values=tensor([1., 2., 3., 4.]), size=(2, 2), nnz=4,
       dtype=torch.float64)
>>> csr.to_dense()
tensor([[1., 2.],
        [3., 4.]], dtype=torch.float64)

Linear Algebra operations（稀疏与稠密之间混合运算）

M表示2-D张量，V表示1-D张量，f表示标量，*表示逐元素乘法，@表示矩阵乘法。M[SparseSemiStructured]表示一种半结构化的稀疏矩阵，此处不再展开，可以自行去torch官网察看。

PyTorch operation	Sparse grad	Layout signature
torch.mv()	no	M[sparse_coo] @ V[strided] -> V[strided]
torch.mv()	no	M[sparse_csr] @ V[strided] -> V[strided]
torch.matmul()	no	M[sparse_coo] @ M[strided] -> M[strided]
torch.matmul()	no	M[sparse_csr] @ M[strided] -> M[strided]
torch.matmul()	no	M[SparseSemiStructured] @ M[strided] -> M[strided]
torch.matmul()	no	M[strided] @ M[SparseSemiStructured] -> M[strided]
torch.mm()	no	M[strided] @ M[SparseSemiStructured] -> M[strided]
torch.mm()	no	M[sparse_coo] @ M[strided] -> M[strided]
torch.mm()	no	M[SparseSemiStructured] @ M[strided] -> M[strided]
torch.sparse.mm()	yes	M[sparse_coo] @ M[strided] -> M[strided]
torch.smm()	no	M[sparse_coo] @ M[strided] -> M[sparse_coo]
torch.hspmm()	no	M[sparse_coo] @ M[strided] -> M[hybrid sparse_coo]
torch.bmm()	no	T[sparse_coo] @ T[strided] -> T[strided]
torch.addmm()	no	f * M[strided] + f * (M[sparse_coo] @ M[strided]) -> M[strided]
torch.addmm()	no	f * M[strided] + f * (M[SparseSemiStructured] @ M[strided]) -> M[strided]
torch.addmm()	no	f * M[strided] + f * (M[strided] @ M[SparseSemiStructured]) -> M[strided]
torch.sparse.addmm()	yes	f * M[strided] + f * (M[sparse_coo] @ M[strided]) -> M[strided]
torch.sspaddmm()	no	f * M[sparse_coo] + f * (M[sparse_coo] @ M[strided]) -> M[sparse_coo]
torch.lobpcg()	no	GENEIG(M[sparse_coo]) -> M[strided], M[strided]
torch.pca_lowrank()	yes	PCA(M[sparse_coo]) -> M[strided], M[strided], M[strided]
torch.svd_lowrank()	yes	SVD(M[sparse_coo]) -> M[strided], M[strided], M[strided]

以上API中，如果 Layout signature中提供了 @或者 *操作符，就不需要记住API，直接通过操作符即可隐式调用对应的API。如：

python 复制代码

>>> a = torch.tensor([[0, 0, 1, 0], [1, 2, 0, 0], [0, 0, 0, 0]], dtype=torch.float64)
>>> sp = a.to_sparse_csr()
>>> vec = torch.randn(4, 1, dtype=torch.float64)
>>> sp.matmul(vec)
tensor([[ 0.4788],
        [-3.2338],
        [ 0.0000]], dtype=torch.float64)
>>> sp @ vec
tensor([[ 0.4788],
        [-3.2338],
        [ 0.0000]], dtype=torch.float64)

需要注意的是，使用操作符在稀疏张量和稠密张量之间乘法运算时，返回的都是稠密张量。如果想要返回稀疏张量，需要显式使用torch.smm()。

torch同样支持稀疏与稀疏之间的运算，但要求输入的稀疏张量必须具有相同的稀疏结构，否则会报错，返回的稀疏张量的稀疏结构也与输入相同。

乘法运算：

python 复制代码

>>> a = torch.tensor([[0, 0, 1, 0], [1, 2, 0, 0], [0, 1, 0, 0], [1, 0, 0, 0]], dtype=torch.float64)
>>> b = torch.tensor([[0, 0, 2, 0], [3, 1, 0, 0], [0, 0, 4, 0], [1, 0, 0, 1]], dtype=torch.float64)
>>> sp1 = a.to_sparse_coo()
>>> sp2 = b.to_sparse_coo()
>>> sp1 @ sp2
tensor(indices=tensor([[0, 1, 1, 1, 2, 2, 3],
                       [2, 0, 1, 2, 0, 1, 2]]),
       values=tensor([4., 6., 2., 2., 3., 1., 2.]),
       size=(4, 4), nnz=7, dtype=torch.float64, layout=torch.sparse_coo)

加法运算

python 复制代码

>>> a = torch.tensor([[0, 0, 1, 0], [1, 2, 0, 0], [0, 1, 0, 0], [1, 0, 0, 0]], dtype=torch.float64)
>>> b = torch.tensor([[0, 0, 2, 0], [3, 1, 0, 0], [0, 0, 4, 0], [1, 0, 0, 1]], dtype=torch.float64)
>>> sp1 = a.to_sparse_coo()
>>> sp2 = b.to_sparse_coo()
>>> sp3 = b.to_sparse_csr()
>>> sp1 + sp2
tensor(indices=tensor([[0, 1, 1, 2, 2, 3, 3],
                       [2, 0, 1, 1, 2, 0, 3]]),
       values=tensor([3., 4., 3., 1., 4., 2., 1.]),
       size=(4, 4), nnz=7, dtype=torch.float64, layout=torch.sparse_coo)
>>> sp1 + sp3
UserWarning: Sparse CSR tensor support is in beta state. If you miss a functionality in the sparse tensor support, please submit a feature request to https://github.com/pytorch/pytorch/issues. (Triggered internally at C:\actions-runner\_work\pytorch\pytorch\builder\windows\pytorch\aten\src\ATen\SparseCsrTensorImpl.cpp:55.)
  sp3 = b.to_sparse_csr()
Traceback (most recent call last):
  File "C:\Users\Xu Han\Desktop\pycharm-projects\MD_notes\main.py", line 18, in <module>
    print(sp1 + sp3)
RuntimeError: memory format option is only supported by strided tensors

Tensor methods and sparse（与稀疏有关的tensor成员函数）

PyTorch operation	return
Tensor.is_sparse	Is`True` if the Tensor uses sparse COO storage layout, `False` otherwise.
Tensor.is_sparse_csr	Is`True` if the Tensor uses sparse CSR storage layout, `False` otherwise.
Tensor.dense_dim	Return the number of dense dimensions in a sparse tensor`self`.
Tensor.sparse_dim	Return the number of sparse dimensions in a sparse tensor`self`.

这里打断一下表格，讲解一下dense_dim和sparse_dim的含义。上文中，我们曾构建过稠密混合的coo张量，如下：

python 复制代码

>>> i = [[0, 1, 1],
         [2, 0, 2]]
>>> v =  [[3, 4], [5, 6], [7, 8]]
>>> s = torch.sparse_coo_tensor(i, v, (2, 3, 2))
>>> s
tensor(indices=tensor([[0, 1, 1],
                       [2, 0, 2]]),
       values=tensor([[3, 4],
                      [5, 6],
                      [7, 8]]),
       size=(2, 3, 2), nnz=3, layout=torch.sparse_coo)

那么，对于这个tensor，它的dense_dim为1，sparse_dim为2。

此外，在进行稀疏与稀疏之间的数学运算时，一定要保证稀疏张量的sparse_dim等于2.

继续表格。

PyTorch operation	return
Tensor.sparse_mask	Returns a new sparse tensor with values from a strided tensor`self` filtered by the indices of the sparse tensor `mask`.
Tensor.to_sparse	Returns a sparse copy of the tensor.
Tensor.to_sparse_coo	Convert a tensor to coordinate format.
Tensor.to_sparse_csr	Convert a tensor to compressed row storage format (CSR).
Tensor.to_sparse_csc	Convert a tensor to compressed column storage (CSC) format.
Tensor.to_sparse_bsr	Convert a tensor to a block sparse row (BSR) storage format of given blocksize.
Tensor.to_sparse_bsc	Convert a tensor to a block sparse column (BSC) storage format of given blocksize.
Tensor.to_dense	Creates a strided copy of`self` if `self` is not a strided tensor, otherwise returns `self`.
Tensor.values	Return the values tensor of a sparse COO tensor.

以下是仅限coo张量的成员：

PyTorch operation	return
Tensor.coalesce	Returns a coalesced copy of`self` if `self` is an uncoalesced tensor.
Tensor.sparse_resize_	Resizes`self` sparse tensor to the desired size and the number of sparse and dense dimensions.
Tensor.sparse_resize_and_clear_	Removes all specified elements from a sparse tensor`self` and resizes `self` to the desired size and the number of sparse and dense dimensions.
Tensor.is_coalesced	Returns`True` if `self` is a sparse COO tensor that is coalesced, `False` otherwise.
Tensor.indices	Return the indices tensor of a sparse COO tensor.

以下是仅限csr和bsr张量的成员：

PyTorch operation	return
Tensor.crow_indices	Returns the tensor containing the compressed row indices of the`self` tensor when `self` is a sparse CSR tensor of layout `sparse_csr`.
Tensor.col_indices	Returns the tensor containing the column indices of the`self` tensor when `self` is a sparse CSR tensor of layout `sparse_csr`.

以下是仅限csc和bsc张量的成员：

PyTorch operation	return
Tensor.row_indices	...
Tensor.ccol_indices	...

coo张量可用的tensor成员函数（经实测，csr也有一些可以用，比如dim()）

add() add_() addmm() addmm_() any() asin() asin_() arcsin() arcsin_() bmm() clone() deg2rad() deg2rad_() detach() detach_() dim() div() div_() floor_divide() floor_divide_() get_device() index_select() isnan() log1p() log1p_() mm() mul() mul_() mv() narrow_copy() neg() neg_() negative() negative_() numel() rad2deg() rad2deg_() resize_as_() size() pow() sqrt() square() smm() sspaddmm() sub() sub_() t() t_() transpose() transpose_() zero_()

Torch functions specific to sparse Tensors（与稀疏有关的torch函数）

PyTorch operation	return
sparse_coo_tensor	Constructs a sparse tensor in COO(rdinate) format with specified values at the given`indices`.
sparse_csr_tensor	Constructs a sparse tensor in CSR (Compressed Sparse Row) with specified values at the given`crow_indices` and `col_indices`.
sparse_csc_tensor	Constructs a sparse tensor in CSC (Compressed Sparse Column) with specified values at the given`ccol_indices` and `row_indices`.
sparse_bsr_tensor	Constructs a sparse tensor in BSR (Block Compressed Sparse Row)) with specified 2-dimensional blocks at the given`crow_indices` and `col_indices`.
sparse_bsc_tensor	Constructs a sparse tensor in BSC (Block Compressed Sparse Column)) with specified 2-dimensional blocks at the given`ccol_indices` and `row_indices`.
sparse_compressed_tensor	Constructs a sparse tensor in Compressed Sparse format - CSR, CSC, BSR, or BSC - with specified values at the given`compressed_indices` and `plain_indices`.
sparse.sum	Return the sum of each row of the given sparse tensor.
sparse.addmm	This function does exact same thing as torch.addmm() in the forward, except that it supports backward for sparse COO matrix`mat1`.
sparse.sampled_addmm	Performs a matrix multiplication of the dense matrices`mat1` and `mat2` at the locations specified by the sparsity pattern of `input`.
sparse.mm	Performs a matrix multiplication of the sparse matrix`mat1`
sspaddmm	Matrix multiplies a sparse tensor`mat1` with a dense tensor `mat2`, then adds the sparse tensor `input` to the result.
hspmm	Performs a matrix multiplication of a sparse COO matrix`mat1` and a strided matrix `mat2`.
smm	Performs a matrix multiplication of the sparse matrix`input` with the dense matrix `mat`.
sparse.softmax	Applies a softmax function.
sparse.log_softmax	Applies a softmax function followed by logarithm.
sparse.spdiags	Creates a sparse 2D tensor by placing the values from rows of`diagonals` along specified diagonals of the output

支持稀疏张量的常规torch函数

cat() dstack() empty() empty_like() hstack() index_select() is_complex() is_floating_point() is_nonzero() is_same_size() is_signed() is_tensor() lobpcg() mm() native_norm() pca_lowrank() select() stack() svd_lowrank() unsqueeze() vstack() zeros() zeros_like()

支持稀疏张量的一元函数

The following operators currently support sparse COO/CSR/CSC/BSR tensor inputs.

abs() asin() asinh() atan() atanh() ceil() conj_physical() floor() log1p() neg() round() sin() sinh() sign() sgn() signbit() tan() tanh() trunc() expm1() sqrt() angle() isinf() isposinf() isneginf() isnan() erf() erfinv()