torch.sparse¶

介绍¶

PyTorch 提供 torch.Tensor 来表示一个包含单一数据类型元素的多维数组。默认情况下，数组元素在内存中是连续存储的，这使得各种依赖于快速访问数组元素的数组处理算法能够高效实现。然而，存在一类重要的多维数组，即所谓的稀疏数组，在这些数组中，数组元素的连续内存存储方式并不理想。稀疏数组的一个特性是其中大部分元素等于零，这意味着如果只存储或处理非零元素，可以节省大量内存和处理器资源。已经开发出多种稀疏存储格式 (如 COO、CSR/CSC、LIL 等。)，这些格式针对稀疏数组中非零元素的特定结构以及对数组的特定操作进行了优化。

注意

当谈到仅存储稀疏数组中的非零元素时，“非零”这一形容词的使用并不是严格的：允许在稀疏数组的数据结构中也存储零。因此，在下文中，我们使用“指定元素”来表示那些实际被存储的数组元素。此外，未指定的元素通常假定其值为零，但并非一定如此，因此我们使用“填充值”这一术语来表示这类元素。

注意

使用稀疏存储格式来存储稀疏数组仅在数组的大小和稀疏程度较高时才具有优势。否则，对于小尺寸或稀疏程度较低的数组，使用连续内存存储格式可能是最有效的方法。

警告

PyTorch 稀疏张量的 API 处于测试阶段，未来可能会发生变化。

稀疏 COO 张量¶

PyTorch 实现了一种所谓的坐标格式（Coordinate format），或 COO 格式，作为实现稀疏张量的存储格式之一。在 COO 格式中，指定的元素以元组的形式存储，其中包含元素的索引和对应值。特别是，

the indices of specified elements are collected in indices tensor of size (ndim, nse) and with element type torch.int64,

the corresponding values are collected in values tensor of size (nse,) and with an arbitrary integer or floating point number element type,

其中 ndim 是张量的维度，nse 是指定元素的数量。

注意

稀疏COO张量的内存消耗至少为 (ndim * 8 + <size of element type in bytes>) * nse 字节（加上存储其他张量数据的固定开销）。

一个步长张量的内存消耗至少为 product(<tensor shape>) * <size of element type in bytes>。

例如，一个10 000 x 10 000张量包含100 000个非零32位浮点数时，使用COO张量布局的内存消耗至少为(2 * 8 + 4) * 100 000 = 2 000 000字节，而使用默认的步长张量布局则为10 000 * 10 000 * 4 = 400 000 000字节。请注意使用COO存储格式所带来的200倍内存节省。

构建¶

可以通过提供两个张量（索引和值）以及稀疏张量的大小（当无法从索引和值张量推断时）来构造稀疏 COO 张量。 torch.sparse_coo_tensor()。

假设我们想要定义一个稀疏张量，在位置 (0, 2) 处的元素为 3，在位置 (1, 0) 处的元素为 4，在位置 (1, 2) 处的元素为 5。未指定的元素默认具有相同的值，即填充值，默认情况下为零。我们可以这样编写：

>>> 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]])

请注意，输入 i 并不是一个索引元组列表。如果你想以这种方式编写索引，应在将其传递给稀疏构造函数之前进行转置：

>>> i = [[0, 2], [1, 0], [1, 2]]
>>> v =  [3,      4,      5    ]
>>> s = torch.sparse_coo_tensor(list(zip(*i)), v, (2, 3))
>>> # Or another equivalent formulation to get s
>>> s = torch.sparse_coo_tensor(torch.tensor(i).t(), v, (2, 3))
>>> torch.sparse_coo_tensor(i.t(), v, torch.Size([2,3])).to_dense()
tensor([[0, 0, 3],
        [4, 0, 5]])

可以通过仅指定其大小来构造一个空的稀疏 COO 张量：

>>> torch.sparse_coo_tensor(size=(2, 3))
tensor(indices=tensor([], size=(2, 0)),
       values=tensor([], size=(0,)),
       size=(2, 3), nnz=0, layout=torch.sparse_coo)

混合稀疏COO张量¶

PyTorch 实现了从标量值稀疏张量到（连续）张量值稀疏张量的扩展。这种张量称为混合张量。

PyTorch 混合 COO 张量通过允许 values 张量为多维张量来扩展稀疏 COO 张量，因此我们有：

the indices of specified elements are collected in indices tensor of size (sparse_dims, nse) and with element type torch.int64,

the corresponding (tensor) values are collected in values tensor of size (nse, dense_dims) and with an arbitrary integer or floating point number element type.

注意

我们使用 (M + K) 维张量来表示一个 N 维混合稀疏张量，其中 M 和 K 分别是稀疏维度和稠密维度的数量，使得 M + K == N 成立。

假设我们想要创建一个 (2 + 1) 维张量，其中位置 (0, 2) 处的条目为 [3, 4]，位置 (1, 0) 处的条目为 [5, 6]，以及位置 (1, 2) 处的条目为 [7, 8]。我们可以这样写

>>> 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)

>>> s.to_dense()
tensor([[[0, 0],
         [0, 0],
         [3, 4]],
        [[5, 6],
         [0, 0],
         [7, 8]]])

一般来说，如果 s 是一个稀疏 COO 张量，并且 M = s.sparse_dim(), K = s.dense_dim()，那么我们有以下不变性：

M + K == len(s.shape) == s.ndim - dimensionality of a tensor is the sum of the number of sparse and dense dimensions,

s.indices().shape == (M, nse) - sparse indices are stored explicitly,

s.values().shape == (nse,) + s.shape[M : M + K] - the values of a hybrid tensor are K-dimensional tensors,

s.values().layout == torch.strided - values are stored as strided tensors.

注意

稠密维度始终跟随稀疏维度，也就是说，不支持稠密和稀疏维度的混合。

未合并的稀疏 COO 张量¶

PyTorch 稀疏 COO 张量格式允许 非合并 稀疏张量，其中索引中可能存在重复坐标；在这种情况下，该索引的值是所有重复值条目的总和。例如，可以为同一个索引 1 指定多个值， 3 和 4，这将导致一个 1-D 非合并张量：

>>> 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.coalesce()
tensor(indices=tensor([[1]]),
       values=tensor([7]),
       size=(3,), nnz=1, layout=torch.sparse_coo)

通常，torch.Tensor.coalesce() 方法的输出是一个具有以下属性的稀疏张量：

指定张量元素的索引是唯一的，
索引按字典顺序排序，
torch.Tensor.is_coalesced() 返回 True。

注意

在大多数情况下，你不需要关心稀疏张量是否已经被合并，因为大多数操作在处理已合并或未合并的稀疏张量时表现是相同的。

然而，某些操作可以在未合并的张量上更高效地实现，而另一些操作则在已合并的张量上更高效。

例如，稀疏 COO 张量的加法是通过简单地连接索引和值张量来实现的：

>>> a = torch.sparse_coo_tensor([[1, 1]], [5, 6], (2,))
>>> b = torch.sparse_coo_tensor([[0, 0]], [7, 8], (2,))
>>> a + b
tensor(indices=tensor([[0, 0, 1, 1]]),
       values=tensor([7, 8, 5, 6]),
       size=(2,), nnz=4, layout=torch.sparse_coo)

如果你反复执行会产生重复条目的操作（例如：torch.Tensor.add()），你应该偶尔合并稀疏张量，以防止它们变得过大。

另一方面，索引的字典序排列对于实现涉及大量元素选择操作（如切片或矩阵乘积）的算法可能是有利的。

使用稀疏 COO 张量¶

让我们考虑以下示例：

>>> i = [[0, 1, 1],
         [2, 0, 2]]
>>> v =  [[3, 4], [5, 6], [7, 8]]
>>> s = torch.sparse_coo_tensor(i, v, (2, 3, 2))

如上所述，稀疏 COO 张量是一个 torch.Tensor 实例，为了将其与使用其他布局的 Tensor 实例区分开来，可以使用 torch.Tensor.is_sparse 或 torch.Tensor.layout 属性：

>>> isinstance(s, torch.Tensor)
True
>>> s.is_sparse
True
>>> s.layout == torch.sparse_coo
True

可以使用方法 torch.Tensor.sparse_dim() 和 torch.Tensor.dense_dim() 分别获取稀疏维度和密集维度的数量。例如：

>>> s.sparse_dim(), s.dense_dim()
(2, 1)

如果 s 是一个稀疏 COO 张量，那么可以使用方法 torch.Tensor.indices() 和 torch.Tensor.values() 获取其 COO 格式的数据。

注意

目前，只有在张量实例被合并的情况下，才能获取 COO 格式的数据：

>>> s.indices()
RuntimeError: Cannot get indices on an uncoalesced tensor, please call .coalesce() first

要获取未合并张量的COO格式数据，请使用 torch.Tensor._values() 和 torch.Tensor._indices()：

>>> s._indices()
tensor([[0, 1, 1],
        [2, 0, 2]])

警告

调用 torch.Tensor._values() 将返回一个 未连接的 张量。要跟踪梯度，必须使用 torch.Tensor.coalesce().values()。

构造一个新的稀疏 COO 张量会生成一个未合并的张量：

>>> s.is_coalesced()
False

但是可以使用 torch.Tensor.coalesce() 方法构建稀疏 COO 张量的合并副本：

>>> s2 = s.coalesce()
>>> s2.indices()
tensor([[0, 1, 1],
       [2, 0, 2]])

在使用未合并的稀疏COO张量时，必须考虑到未合并数据的加法性质：相同索引的值是求和项，其计算结果给出对应张量元素的值。例如，对未合并的稀疏张量进行标量乘法可以通过将所有未合并的值与标量相乘来实现，因为 c * (a + b) == c * a + c * b 成立。然而，任何非线性操作，比如平方根，都无法通过对未合并的数据应用该操作来实现，因为在一般情况下 sqrt(a + b) == sqrt(a) + sqrt(b) 不成立。

仅支持对稀疏 COO 张量的密集维度进行（正步长）切片。索引操作同时支持稀疏和密集维度：

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

在 PyTorch 中，稀疏张量的填充值无法显式指定，通常假定为零。然而，存在一些操作可能会以不同的方式解释填充值。例如，torch.sparse.softmax() 在假设填充值为负无穷的情况下计算 softmax。

稀疏CSR张量¶

CSR（压缩稀疏行）稀疏张量格式实现了用于存储二维张量的CSR格式。虽然目前不支持N维张量，但与COO格式相比，其主要优势在于更高效地利用存储空间，并且使用MKL和MAGMA后端时，可以显著加快诸如稀疏矩阵-向量乘法等计算操作。截至目前，尚不支持CUDA。

一个CSR稀疏张量由三个1-D张量组成：crow_indices、col_indices 和 values：

The crow_indices tensor consists of compressed row indices. This is a 1-D tensor of size size[0] + 1. The last element is the number of non-zeros. This tensor encodes the index in values and col_indices depending on where the given row starts. Each successive number in the tensor subtracted by the number before it denotes the number of elements in a given row.

The col_indices tensor contains the column indices of each value. This is a 1-D tensor of size nnz.

The values tensor contains the values of the CSR tensor. This is a 1-D tensor of size nnz.

注意

索引张量 crow_indices 和 col_indices 的元素类型应为 torch.int64（默认）或 torch.int32。如果你想使用启用 MKL 的矩阵运算，请使用 torch.int32。这是由于 pytorch 默认链接的是 MKL LP64，它使用 32 位整数索引。

CSR 张量的构建¶

稀疏CSR矩阵可以通过使用 torch.sparse_csr_tensor() 方法直接构造。用户必须分别提供行和列的索引以及值张量。 size 参数是可选的，如果未提供，将从 crow_indices 和 col_indices 推断出来。

>>> 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.double)
>>> 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)

CSR张量操作¶

从一个strided或稀疏COO张量构造稀疏CSR张量的最简单方法是使用tensor.to_sparse_csr()。在(strided)张量中的任何零值都将被解释为稀疏张量中的缺失值:

>>> a = torch.tensor([[0, 0, 1, 0], [1, 2, 0, 0], [0, 0, 0, 0]], dtype = torch.float64)
>>> sp = a.to_sparse_csr()
>>> sp
tensor(crow_indices=tensor([0, 1, 3, 3]),
      col_indices=tensor([2, 0, 1]),
      values=tensor([1., 1., 2.]), size=(3, 4), nnz=3, dtype=torch.float64)

稀疏矩阵-向量乘法可以通过tensor.matmul()方法执行。目前这是CSR张量上唯一支持的数学运算。

>>> vec = torch.randn(4, 1, dtype=torch.float64)
>>> sp.matmul(vec)
tensor([[0.9078],
        [1.3180],
        [0.0000]], dtype=torch.float64)

支持的线性代数运算¶

下表总结了稀疏矩阵上的线性代数运算，其中操作数的布局可能不同。此处 T[layout] 表示具有给定布局的张量。同样， M[layout] 表示一个矩阵（2-D PyTorch 张量），而 V[layout] 表示一个向量（1-D PyTorch 张量）。此外，f 表示一个标量（浮点数或 0-D PyTorch 张量），* 是逐元素乘法，而 @ 是矩阵乘法。

PyTorch 操作	稀疏梯度？	布局签名
`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.mm()`	no	`M[sparse_coo] @ M[strided] -> M[strided]`
`torch.sparse.mm()`	是的	`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.sparse.addmm()`	是的	`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()`	是的	`PCA(M[sparse_coo]) -> M[strided], M[strided], M[strided]`
`torch.svd_lowrank()`	是的	`SVD(M[sparse_coo]) -> M[strided], M[strided], M[strided]`

其中“是否稀疏梯度？”列表示 PyTorch 操作是否支持关于稀疏矩阵参数的反向传播。除了 torch.smm()，所有 PyTorch 操作都支持关于步距矩阵参数的反向传播。

注意

目前，PyTorch 不支持具有布局签名 M[strided] @ M[sparse_coo] 的矩阵乘法。但是，应用程序仍可以使用矩阵关系 D @ S == (S.t() @ D.t()).t() 来计算此操作。

张量方法和稀疏¶

以下 Tensor 方法与稀疏张量相关：

`Tensor.is_sparse`	如果是使用稀疏存储布局的Tensor，则为 `True`，否则为 `False`。
`Tensor.dense_dim`	返回稀疏张量 `self`中的密集维度数量。
`Tensor.sparse_dim`	返回稀疏张量中的稀疏维度数量稀疏张量 `self`。
`Tensor.sparse_mask`	返回一个新的稀疏张量，其值来自步长张量`self`，并通过稀疏张量`mask`的索引进行过滤。
`Tensor.to_sparse`	返回张量的稀疏副本。
`Tensor.to_sparse_coo`	将张量转换为坐标格式。
`Tensor.to_sparse_csr`	将张量转换为压缩行存储格式。
`Tensor.indices`	返回稀疏COO张量的索引张量。
`Tensor.values`	返回稀疏COO张量的值张量。

以下 Tensor 方法特定于稀疏 COO 张量：

`Tensor.coalesce`	如果 `self` 是一个非合并张量，则返回 `self` 的合并副本。
`Tensor.sparse_resize_`	调整 `self` 稀疏张量到所需的大小和稀疏及稠密维度的数量。
`Tensor.sparse_resize_and_clear_`	从指定的稀疏张量 `self` 中删除所有指定元素，并将 `self` 调整为所需的大小和稀疏及稠密维度的数量。
`Tensor.is_coalesced`	如果 `self` 是一个合并的稀疏COO张量，则返回 `True`；否则返回 `False`。
`Tensor.to_dense`	创建一个步长为 `self` 的副本。

以下方法是特定于稀疏CSR张量的：

`Tensor.crow_indices`	返回包含`self`张量的压缩行索引的张量，当`self`是布局为`sparse_csr`的稀疏CSR张量时。
`Tensor.col_indices`	返回包含`self`张量的列索引的张量，当`self`是布局为`sparse_csr`的稀疏CSR张量时。

以下 Tensor 方法支持稀疏 COO 张量：

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_()

与稀疏张量相关的PyTorch函数¶

`sparse_coo_tensor`	使用指定的值在给定的`indices`处构造一个COO（坐标）格式的稀疏张量。
`sparse_csr_tensor`	使用指定的值在给定的`crow_indices`和`col_indices`处构造一个CSR（压缩稀疏行）格式的稀疏张量。
`sparse.sum`	返回稀疏张量 `input` 在给定维度 `dim` 中每一行的总和。
`sparse.addmm`	此函数在前向传播中的作用与 `torch.addmm()` 完全相同，不同之处在于它支持稀疏 COO 矩阵 `mat1` 的反向传播。
`sparse.sampled_addmm`	执行密集矩阵`mat1`和`mat2`的矩阵乘法，位置由稀疏模式指定的`input`。
`sparse.mm`	对稀疏矩阵 `mat1` 执行矩阵乘法
`sspaddmm`	矩阵将稀疏张量 `mat1` 与密集张量 `mat2` 相乘，然后将稀疏张量 `input` 加到结果中。
`hspmm`	执行一个稀疏COO矩阵 `mat1` 和一个步幅矩阵 `mat2` 的矩阵乘法。
`smm`	对稀疏矩阵 `input` 与密集矩阵 `mat` 执行矩阵乘法。
`sparse.softmax`	应用一个 softmax 函数。
`sparse.log_softmax`	应用一个 softmax 函数，然后应用对数函数。

其他函数¶

以下 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()