目录
快捷方式

torch.sparse

警告

稀疏张量的 PyTorch API 处于测试阶段，未来可能会发生变化。我们热烈欢迎将功能请求、错误报告和一般建议作为 GitHub 问题提出。

为什么以及何时使用稀疏性

默认情况下，PyTorch 在物理内存中连续存储 torch.Tensor 元素。这使得各种数组处理算法的实现非常高效，这些算法需要快速访问元素。

现在，一些用户可能会决定使用张量来表示数据，例如图邻接矩阵、修剪权重或点云，这些张量的元素大多数为零值。我们认识到这些重要的应用，并旨在通过稀疏存储格式为这些用例提供性能优化。

近年来，已经开发出各种稀疏存储格式，例如 COO、CSR/CSC、半结构化、LIL 等。虽然它们在精确布局上有所不同，但它们都通过对零值元素的有效表示来压缩数据。我们将未压缩的值称为指定，与未指定的压缩元素形成对比。

通过压缩重复零，稀疏存储格式旨在节省各种 CPU 和 GPU 上的内存和计算资源。特别是对于高度稀疏性或高度结构化稀疏性，这可能会对性能产生重大影响。因此，稀疏存储格式可以看作是性能优化。

与许多其他性能优化一样，稀疏存储格式并不总是优势。在尝试将稀疏格式用于您的用例时，您可能会发现执行时间增加而不是减少。

如果您从分析上预计会看到性能的显著提升，但实际测量到性能下降，请随时在 GitHub 上提 issue。这将有助于我们优先考虑高效内核的实现和更广泛的性能优化。

我们简化了尝试不同稀疏性布局并进行转换的过程，而不必对最适合您的特定应用程序的格式有偏见。

功能概述

我们希望能够通过为每种布局提供转换例程，从给定的密集张量直接构建稀疏张量。

在下面的示例中，我们将具有默认密集（带步长）布局的二维张量转换为由 COO 内存布局支持的二维张量。在这种情况下，只存储非零元素的值和索引。

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

PyTorch 目前支持 COO、CSR、CSC、BSR 和 BSC。

我们还有一个原型实现，用于支持：ref: 半结构化稀疏性<sparse-semi-structured-docs>。有关更多详细信息，请参阅参考资料。

请注意，我们对这些格式进行了细微的泛化。

批处理：诸如 GPU 之类的设备需要批处理才能获得最佳性能，因此我们支持批处理维度。

我们目前提供了一个非常简单的批处理版本，其中稀疏格式的每个组件本身都进行批处理。这也要求每个批处理条目具有相同数量的指定元素。在此示例中，我们从三维密集张量构建三维（批处理）CSR 张量。

>>> t = torch.tensor([[[1., 0], [2., 3.]], [[4., 0], [5., 6.]]])
>>> t.dim()
3
>>> t.to_sparse_csr()
tensor(crow_indices=tensor([[0, 1, 3],
                            [0, 1, 3]]),
       col_indices=tensor([[0, 0, 1],
                           [0, 0, 1]]),
       values=tensor([[1., 2., 3.],
                      [4., 5., 6.]]), size=(2, 2, 2), nnz=3,
       layout=torch.sparse_csr)

密集维度：另一方面，某些数据（如图嵌入）可能最好被视为稀疏向量集合，而不是标量集合。

在此示例中，我们从三维带步长张量创建具有 2 个稀疏维度和 1 个密集维度的三维混合 COO 张量。如果三维带步长张量中的整行都为零，则不进行存储。但是，如果该行中的任何值非零，则将全部存储这些值。这减少了索引数量，因为我们每行只需要一个索引，而不是每个元素都需要一个索引。但它也增加了值的存储量。因为只有完全为零的行可以被忽略，而任何非零值的出现都会导致整行被存储。

>>> t = torch.tensor([[[0., 0], [1., 2.]], [[0., 0], [3., 4.]]])
>>> t.to_sparse(sparse_dim=2)
tensor(indices=tensor([[0, 1],
                       [1, 1]]),
       values=tensor([[1., 2.],
                      [3., 4.]]),
       size=(2, 2, 2), nnz=2, layout=torch.sparse_coo)

运算符概述

从根本上说，对具有稀疏存储格式的张量的运算与对具有带步长（或其他）存储格式的张量的运算相同。存储的具体细节，即数据的物理布局，会影响运算的性能，但不应影响语义。

我们正在积极增加对稀疏张量的运算符覆盖范围。用户不应期望与密集张量相同的支持级别。有关列表，请参阅我们的 运算符 文档。

>>> b = torch.tensor([[0, 0, 1, 2, 3, 0], [4, 5, 0, 6, 0, 0]])
>>> b_s = b.to_sparse_csr()
>>> b_s.cos()
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
RuntimeError: unsupported tensor layout: SparseCsr
>>> b_s.sin()
tensor(crow_indices=tensor([0, 3, 6]),
       col_indices=tensor([2, 3, 4, 0, 1, 3]),
       values=tensor([ 0.8415,  0.9093,  0.1411, -0.7568, -0.9589, -0.2794]),
       size=(2, 6), nnz=6, layout=torch.sparse_csr)

如上例所示，我们不支持非零保持的单目运算符，例如 cos。非零保持的单目运算符的输出将无法像输入一样充分利用稀疏存储格式，并可能导致内存的灾难性增加。相反，我们依赖用户首先显式转换为密集张量，然后执行运算。

>>> b_s.to_dense().cos()
tensor([[ 1.0000, -0.4161],
        [-0.9900,  1.0000]])

我们知道一些用户希望在进行 cos 等运算时忽略压缩零，而不是保留运算的精确语义。为此，我们可以指向 torch.masked 及其 MaskedTensor，它反过来也由稀疏存储格式和内核支持和驱动。

另外请注意，目前，用户无法选择输出布局。例如，将稀疏张量添加到常规带步长张量会导致带步长张量。一些用户可能更希望这保持稀疏布局，因为他们知道结果仍然会足够稀疏。

>>> a + b.to_sparse()
tensor([[0., 3.],
        [3., 0.]])

我们承认，访问能够有效地生成不同输出布局的内核非常有用。后续操作可能会从接收特定布局中受益匪浅。我们正在开发一个 API 来控制结果布局，并认识到这是一个重要的功能，可以为任何给定模型规划更优化的执行路径。

稀疏半结构化张量

警告

稀疏半结构化张量目前是一个原型功能，可能会发生变化。如果您遇到任何错误，或者想要分享反馈，请随时提交问题。

半结构化稀疏性是一种稀疏数据布局，最初在 NVIDIA 的 Ampere 架构中引入。它也被称为 **细粒度结构化稀疏性** 或 **2:4 结构化稀疏性**。

此稀疏布局将存储每 2n 个元素中的 n 个元素，其中 n 由张量数据类型 (dtype) 的宽度决定。最常用的 dtype 是 float16，其中 n=2，因此被称为“2:4 结构化稀疏性”。

在 这篇 NVIDIA 博客文章 中对半结构化稀疏性进行了更详细的解释。

在 PyTorch 中，半结构化稀疏性是通过一个张量子类实现的。通过子类化，我们可以覆盖 __torch_dispatch__，允许我们在执行矩阵乘法时使用更快的稀疏内核。我们还可以将张量以其压缩形式存储在子类中，以减少内存开销。

在这种压缩形式中，稀疏张量通过仅保留 *指定* 的元素和一些元数据（编码掩码）来存储。

注意

半结构化稀疏张量的指定元素和元数据掩码一起存储在一个扁平的压缩张量中。它们彼此追加，形成一个连续的内存块。

compressed tensor = [ 原始张量的指定元素 | metadata_mask ]

对于大小为 (r, c) 的原始张量，我们预期前 m * k // 2 个元素是保留的元素，其余的张量是元数据。

为了方便用户查看指定元素和掩码，可以使用 .indices() 和 .values() 分别访问掩码和指定元素。

  • .values() 返回一个大小为 (r, c//2) 且与密集矩阵具有相同 dtype 的张量中的指定元素。

  • .indices() 返回一个大小为 (r, c//2 ) 的张量中的 metadata_mask，如果 dtype 是 torch.float16 或 torch.bfloat16，则元素类型为 torch.int16，如果 dtype 是 torch.int8，则元素类型为 torch.int32。

对于 2:4 稀疏张量，元数据开销很小 - 每个指定元素仅 2 位。

注意

重要的是要注意，torch.float32 仅支持 1:2 稀疏性。因此，它不遵循与上述相同的公式。

这里，我们分解了如何计算 2:4 稀疏张量的压缩率（密集大小 / 稀疏大小）。

令 (r, c) = tensor.shape 且 e = bitwidth(tensor.dtype)，因此对于 torch.float16 和 torch.bfloat16，e = 16，对于 torch.int8，e = 8。

Mdense=r×c×eMsparse=Mspecified+Mmetadata=r×c2×e+r×c2×2=rce2+rc=rce(12+1e)M_{dense} = r \times c \times e \\ M_{sparse} = M_{specified} + M_{metadata} = r \times \frac{c}{2} \times e + r \times \frac{c}{2} \times 2 = \frac{rce}{2} + rc =rce(\frac{1}{2} +\frac{1}{e})

利用这些计算，我们可以确定原始密集表示和新的稀疏表示的总内存占用。

这为我们提供了一个简单的压缩率公式，该公式仅取决于张量数据类型的位宽。

C=MsparseMdense=12+1eC = \frac{M_{sparse}}{M_{dense}} = \frac{1}{2} + \frac{1}{e}

通过使用此公式，我们发现对于 torch.float16 或 torch.bfloat16 压缩率为 56.25%，而对于 torch.int8 压缩率为 62.5%。

构建稀疏半结构化张量

您可以通过简单地使用 torch.to_sparse_semi_structured 函数将密集张量转换为稀疏半结构化张量。

请注意，我们只支持 CUDA 张量，因为半结构化稀疏性的硬件兼容性仅限于 NVIDIA GPU。

以下数据类型支持半结构化稀疏性。请注意，每个数据类型都有自己的形状约束和压缩因子。

PyTorch 数据类型

形状约束

压缩因子

稀疏性模式

torch.float16

张量必须为 2D，(r, c) 必须都是 64 的正倍数

9/16

2:4

torch.bfloat16

张量必须为 2D，(r, c) 必须都是 64 的正倍数

9/16

2:4

torch.int8

张量必须为 2D，(r, c) 必须都是 128 的正倍数

10/16

2:4

要构建一个半结构化稀疏张量，首先创建一个符合 2:4（或半结构化）稀疏格式的普通密集张量。为此，我们平铺一个小的 1x4 条带以创建一个 16x16 密集 float16 张量。之后，我们可以调用 to_sparse_semi_structured 函数来压缩它以进行加速推理。

>>> from torch.sparse import to_sparse_semi_structured
>>> A = torch.Tensor([0, 0, 1, 1]).tile((128, 32)).half().cuda()
tensor([[0., 0., 1.,  ..., 0., 1., 1.],
        [0., 0., 1.,  ..., 0., 1., 1.],
        [0., 0., 1.,  ..., 0., 1., 1.],
        ...,
        [0., 0., 1.,  ..., 0., 1., 1.],
        [0., 0., 1.,  ..., 0., 1., 1.],
        [0., 0., 1.,  ..., 0., 1., 1.]], device='cuda:0', dtype=torch.float16)
>>> A_sparse = to_sparse_semi_structured(A)
SparseSemiStructuredTensor(shape=torch.Size([128, 128]), transposed=False, values=tensor([[1., 1., 1.,  ..., 1., 1., 1.],
        [1., 1., 1.,  ..., 1., 1., 1.],
        [1., 1., 1.,  ..., 1., 1., 1.],
        ...,
        [1., 1., 1.,  ..., 1., 1., 1.],
        [1., 1., 1.,  ..., 1., 1., 1.],
        [1., 1., 1.,  ..., 1., 1., 1.]], device='cuda:0', dtype=torch.float16), metadata=tensor([[-4370, -4370, -4370,  ..., -4370, -4370, -4370],
        [-4370, -4370, -4370,  ..., -4370, -4370, -4370],
        [-4370, -4370, -4370,  ..., -4370, -4370, -4370],
        ...,
        [-4370, -4370, -4370,  ..., -4370, -4370, -4370],
        [-4370, -4370, -4370,  ..., -4370, -4370, -4370],
        [-4370, -4370, -4370,  ..., -4370, -4370, -4370]], device='cuda:0',
dtype=torch.int16))

稀疏半结构化张量运算

目前，以下运算支持半结构化稀疏张量

  • torch.addmm(bias, dense, sparse.t())

  • torch.mm(dense, sparse)

  • torch.mm(sparse, dense)

  • aten.linear.default(dense, sparse, bias)

  • aten.t.default(sparse)

  • aten.t.detach(sparse)

要使用这些运算，只需传递 to_sparse_semi_structured(tensor) 的输出，而不要使用 tensor，一旦您的张量以半结构化稀疏格式包含 0，就像这样

>>> a = torch.Tensor([0, 0, 1, 1]).tile((64, 16)).half().cuda()
>>> b = torch.rand(64, 64).half().cuda()
>>> c = torch.mm(a, b)
>>> a_sparse = to_sparse_semi_structured(a)
>>> torch.allclose(c, torch.mm(a_sparse, b))
True

使用半结构化稀疏性加速 nn.Linear

如果权重已经是半结构化稀疏的，您可以通过几行代码加速模型中的线性层

>>> input = torch.rand(64, 64).half().cuda()
>>> mask = torch.Tensor([0, 0, 1, 1]).tile((64, 16)).cuda().bool()
>>> linear = nn.Linear(64, 64).half().cuda()
>>> linear.weight = nn.Parameter(to_sparse_semi_structured(linear.weight.masked_fill(~mask, 0)))

稀疏 COO 张量

PyTorch 实现所谓的坐标格式，或 COO 格式，作为实现稀疏张量的存储格式之一。在 COO 格式中，指定的元素存储为元素索引和相应值的元组。特别是，

  • 指定元素的索引在大小为 (ndim, nse) 的 indices 张量中收集，元素类型为 torch.int64，

  • 相应的数值在大小为 (nse,) 的 values 张量中收集，元素类型为任意整数或浮点数，

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

注意

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

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

例如，包含 100 000 个非零 32 位浮点数的 10 000 x 10 000 张量的内存消耗至少为 (2 * 8 + 4) * 100 000 = 2 000 000 字节，使用 COO 张量布局，而使用默认带步长张量布局则为 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 张量，因此我们有

  • 指定元素的索引在大小为 (sparse_dims, nse) 的 indices 张量中收集，元素类型为 torch.int64，

  • 相应的（张量）数值在大小为 (nse, dense_dims) 的 values 张量中收集，元素类型为任意整数或浮点数。

注意

我们使用 (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 - 张量的维数是稀疏维数和密集维数之和，

  • s.indices().shape == (M, nse) - 稀疏索引被显式存储，

  • s.values().shape == (nse,) + s.shape[M : M + K] - 混合张量的值是 K 维张量。

  • s.values().layout == torch.strided - 值存储为步长张量。

注意

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

注意

为了确保构建的稀疏张量具有一致的索引、值和大小，可以通过 check_invariants=True 关键字参数在每个张量创建时启用不变性检查，或者使用 torch.sparse.check_sparse_tensor_invariants 上下文管理器实例全局启用。默认情况下，稀疏张量不变性检查处于禁用状态。

非合并稀疏 COO 张量

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

>>> 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() 方法的输出是一个稀疏张量，具有以下属性

注意

在大多数情况下，你不需要关心稀疏张量是否合并，因为大多数操作在给定合并或非合并的稀疏张量时会以相同的方式工作。

但是，某些操作可以在非合并张量上更有效地实现，而另一些操作则可以在合并张量上更有效地实现。

例如，稀疏 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) 格式，PyTorch 稀疏压缩张量通过支持稀疏张量批次、允许多维张量值以及将稀疏张量值存储在稠密块中来扩展此格式。

注意

我们使用 (B + M + K) 维张量来表示一个 N 维稀疏压缩混合张量，其中 B、M 和 K 分别是批次、稀疏和稠密维度的数量，使得 B + M + K == N 成立。稀疏压缩张量的稀疏维度数量始终为 2，M == 2。

注意

我们说索引张量 compressed_indices 使用 CSR 编码，如果满足以下不变性

  • compressed_indices 是一个连续的步长 32 位或 64 位整数张量

  • compressed_indices 形状为 (*batchsize, compressed_dim_size + 1)，其中 compressed_dim_size 是压缩维度的数量（例如行或列）

  • compressed_indices[..., 0] == 0，其中 ... 表示批次索引

  • compressed_indices[..., compressed_dim_size] == nse，其中 nse 是指定元素的数量

  • 0 <= compressed_indices[..., i] - compressed_indices[..., i - 1] <= plain_dim_size 对于 i=1, ..., compressed_dim_size，其中 plain_dim_size 是普通维度的数量（正交于压缩维度，例如列或行）。

为了确保构建的稀疏张量具有一致的索引、值和大小，可以通过 check_invariants=True 关键字参数在每个张量创建时启用不变性检查，或者使用 torch.sparse.check_sparse_tensor_invariants 上下文管理器实例全局启用。默认情况下，稀疏张量不变性检查处于禁用状态。

注意

将稀疏压缩布局推广到 N 维张量会导致有关指定元素计数的一些混淆。当稀疏压缩张量包含批次维度时，指定元素的数量将对应于每个批次中此类元素的数量。当稀疏压缩张量具有稠密维度时，所考虑的元素现在是 K 维数组。同样，对于块稀疏压缩布局，2D 块被视为被指定的元素。例如，一个 3 维块稀疏张量，它具有一个长度为 b 的批次维度，以及一个块形状为 p, q。如果此张量有 n 个指定元素，那么实际上我们每个批次有 n 个指定的块。此张量将具有形状为 (b, n, p, q) 的 values。这种对指定元素数量的解释源于所有稀疏压缩布局都源于 2 维矩阵的压缩。批次维度被视为稀疏矩阵的堆叠，稠密维度改变了元素的含义，从简单的标量值到具有自身维度的数组。

稀疏 CSR 张量

与 COO 格式相比，CSR 格式的主要优势是更好地利用存储空间以及使用 MKL 和 MAGMA 后端进行稀疏矩阵向量乘法的速度快得多。

在最简单的情况下，一个 (0 + 2 + 0) 维稀疏 CSR 张量由三个 1 维张量组成：crow_indices、col_indices 和 values

  • 张量 crow_indices 包含压缩的行索引。这是一个大小为 nrows + 1（行数加 1）的 1 维张量。 crow_indices 的最后一个元素是指定元素的数量 nse。该张量根据给定行开始的位置，编码了 values 和 col_indices 中的索引。张量中每个连续的数字减去之前的数字，表示给定行中元素的数量。

  • 张量 col_indices 包含每个元素的列索引。这是一个大小为 nse 的 1 维张量。

  • 张量 values 包含 CSR 张量元素的值。这是一个大小为 nse 的 1 维张量。

注意

索引张量 crow_indices 和 col_indices 的元素类型应该为 torch.int64（默认）或 torch.int32。如果要使用支持 MKL 的矩阵运算，请使用 torch.int32。这是因为 PyTorch 默认情况下链接到 MKL LP64，该库使用 32 位整数索引。

在一般情况下，(B + 2 + K) 维稀疏 CSR 张量由两个 (B + 1) 维索引张量 crow_indices 和 col_indices 以及一个 (1 + K) 维 values 张量组成，使得

  • crow_indices.shape == (*batchsize, nrows + 1)

  • col_indices.shape == (*batchsize, nse)

  • values.shape == (nse, *densesize)

而稀疏 CSR 张量的形状为 (*batchsize, nrows, ncols, *densesize)，其中 len(batchsize) == B 且 len(densesize) == K。

注意

稀疏 CSR 张量的批次是相互依赖的：所有批次中指定元素的数量必须相同。这个有些人为的限制允许有效地存储不同 CSR 批次的索引。

注意

可以使用 torch.Tensor.sparse_dim() 和 torch.Tensor.dense_dim() 方法获取稀疏和密集维度的数量。批次维度可以从张量形状计算得出：batchsize = tensor.shape[:-tensor.sparse_dim() - tensor.dense_dim()]。

注意

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

使用与 稀疏 COO 格式介绍中的注释 相同的示例数据，一个 10000 x 10000 的张量，包含 100000 个非零的 32 位浮点数，使用 CSR 张量布局，其内存消耗至少为 (10000 * 8 + (8 + 4 * 1) * 100 000) * 1 = 1 280 000 字节。请注意，与使用 COO 和跨步格式相比，使用 CSR 存储格式分别节省了 1.6 倍和 310 倍的内存。

CSR 张量的构造

可以使用 torch.sparse_csr_tensor() 函数直接构造稀疏 CSR 张量。用户必须分别提供行和列索引以及值张量，其中行索引必须使用 CSR 压缩编码指定。 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.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)

注意

推断出的 size 中稀疏维度的值，从 crow_indices 的大小和 col_indices 中的最大索引值计算得出。如果列数需要大于推断出的 size 中的值，则必须显式指定 size 参数。

从跨步张量或稀疏 COO 张量构造二维稀疏 CSR 张量最简单的方法是使用 torch.Tensor.to_sparse_csr() 方法。 (跨步) 张量中的任何零值都将被解释为稀疏张量中的缺失值。

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

CSR 张量运算

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

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

稀疏 CSC 张量

稀疏 CSC（压缩稀疏列）张量格式实现了 CSC 格式，用于存储二维张量，并扩展为支持稀疏 CSC 张量的批次，以及值为多维张量。

注意

当转置操作是交换稀疏维度时，稀疏 CSC 张量本质上是稀疏 CSR 张量的转置。

与 稀疏 CSR 张量 类似，稀疏 CSC 张量由三个张量组成：ccol_indices、row_indices 和 values

  • 张量 ccol_indices 包含压缩的列索引。这是一个形状为 (*batchsize, ncols + 1) 的 (B + 1) 维张量。最后一个元素是指定元素的数量 nse。该张量根据给定列开始的位置，编码了 values 和 row_indices 中的索引。张量中每个连续的数字减去之前的数字，表示给定列中元素的数量。

  • 张量 row_indices 包含每个元素的行索引。这是一个形状为 (*batchsize, nse) 的 (B + 1) 维张量。

  • 张量 values 包含 CSC 张量元素的值。这是一个形状为 (nse, *densesize) 的 (1 + K) 维张量。

CSC 张量的构造

可以使用 torch.sparse_csc_tensor() 函数直接构造稀疏 CSC 张量。用户必须分别提供行和列索引以及值张量，其中列索引必须使用 CSR 压缩编码指定。 size 参数是可选的，如果不存在，则会从 row_indices 和 ccol_indices 张量中推断出来。

>>> ccol_indices = torch.tensor([0, 2, 4])
>>> row_indices = torch.tensor([0, 1, 0, 1])
>>> values = torch.tensor([1, 2, 3, 4])
>>> csc = torch.sparse_csc_tensor(ccol_indices, row_indices, values, dtype=torch.float64)
>>> csc
tensor(ccol_indices=tensor([0, 2, 4]),
       row_indices=tensor([0, 1, 0, 1]),
       values=tensor([1., 2., 3., 4.]), size=(2, 2), nnz=4,
       dtype=torch.float64, layout=torch.sparse_csc)
>>> csc.to_dense()
tensor([[1., 3.],
        [2., 4.]], dtype=torch.float64)

注意

稀疏 CSC 张量构造函数的压缩列索引参数位于行索引参数之前。

可以使用 torch.Tensor.to_sparse_csc() 方法从任何二维张量构造 (0 + 2 + 0) 维稀疏 CSC 张量。 (跨步) 张量中的任何零值都将被解释为稀疏张量中的缺失值。

>>> a = torch.tensor([[0, 0, 1, 0], [1, 2, 0, 0], [0, 0, 0, 0]], dtype=torch.float64)
>>> sp = a.to_sparse_csc()
>>> sp
tensor(ccol_indices=tensor([0, 1, 2, 3, 3]),
       row_indices=tensor([1, 1, 0]),
       values=tensor([1., 2., 1.]), size=(3, 4), nnz=3, dtype=torch.float64,
       layout=torch.sparse_csc)

稀疏 BSR 张量

稀疏 BSR（块压缩稀疏行）张量格式实现了 BSR 格式，用于存储二维张量，并扩展为支持稀疏 BSR 张量的批次，以及值为多维张量的块。

稀疏 BSR 张量由三个张量组成：crow_indices、col_indices 和 values

  • 张量 crow_indices 包含压缩的行索引。这是一个形状为 (*batchsize, nrowblocks + 1) 的 (B + 1) 维张量。最后一个元素是指定块的数量 nse。该张量根据给定列块开始的位置，编码了 values 和 col_indices 中的索引。张量中每个连续的数字减去之前的数字，表示给定行中块的数量。

  • 张量 col_indices 包含每个元素的列块索引。这是一个形状为 (*batchsize, nse) 的 (B + 1) 维张量。

  • The values tensor contains the values of the sparse BSR tensor elements collected into two-dimensional blocks. This is a (1 + 2 + K)-D tensor of shape (nse, nrowblocks, ncolblocks, *densesize).

BSR 张量的构造

可以使用 torch.sparse_bsr_tensor() 函数直接构造稀疏 BSR 张量。用户必须分别提供行和列块索引以及值张量，其中行块索引必须使用 CSR 压缩编码指定。如果 size 参数不存在，则会从 crow_indices 和 col_indices 张量中推断出来。

>>> crow_indices = torch.tensor([0, 2, 4])
>>> col_indices = torch.tensor([0, 1, 0, 1])
>>> values = torch.tensor([[[0, 1, 2], [6, 7, 8]],
...                        [[3, 4, 5], [9, 10, 11]],
...                        [[12, 13, 14], [18, 19, 20]],
...                        [[15, 16, 17], [21, 22, 23]]])
>>> bsr = torch.sparse_bsr_tensor(crow_indices, col_indices, values, dtype=torch.float64)
>>> bsr
tensor(crow_indices=tensor([0, 2, 4]),
       col_indices=tensor([0, 1, 0, 1]),
       values=tensor([[[ 0.,  1.,  2.],
                       [ 6.,  7.,  8.]],
                      [[ 3.,  4.,  5.],
                       [ 9., 10., 11.]],
                      [[12., 13., 14.],
                       [18., 19., 20.]],
                      [[15., 16., 17.],
                       [21., 22., 23.]]]),
       size=(4, 6), nnz=4, dtype=torch.float64, layout=torch.sparse_bsr)
>>> bsr.to_dense()
tensor([[ 0.,  1.,  2.,  3.,  4.,  5.],
        [ 6.,  7.,  8.,  9., 10., 11.],
        [12., 13., 14., 15., 16., 17.],
        [18., 19., 20., 21., 22., 23.]], dtype=torch.float64)

可以使用 torch.Tensor.to_sparse_bsr() 方法从任何二维张量构造 (0 + 2 + 0) 维稀疏 BSR 张量，该方法还需要指定值块大小。

>>> dense = torch.tensor([[0, 1, 2, 3, 4, 5],
...                       [6, 7, 8, 9, 10, 11],
...                       [12, 13, 14, 15, 16, 17],
...                       [18, 19, 20, 21, 22, 23]])
>>> bsr = dense.to_sparse_bsr(blocksize=(2, 3))
>>> bsr
tensor(crow_indices=tensor([0, 2, 4]),
       col_indices=tensor([0, 1, 0, 1]),
       values=tensor([[[ 0,  1,  2],
                       [ 6,  7,  8]],
                      [[ 3,  4,  5],
                       [ 9, 10, 11]],
                      [[12, 13, 14],
                       [18, 19, 20]],
                      [[15, 16, 17],
                       [21, 22, 23]]]), size=(4, 6), nnz=4,
       layout=torch.sparse_bsr)

稀疏 BSC 张量

稀疏 BSC（块压缩稀疏列）张量格式实现了 BSC 格式，用于存储二维张量，并扩展到支持稀疏 BSC 张量批处理，以及将值作为多维张量的块。

稀疏 BSC 张量由三个张量组成：ccol_indices、row_indices 和 values

  • The ccol_indices tensor consists of compressed column indices. This is a (B + 1)-D tensor of shape (*batchsize, ncolblocks + 1). The last element is the number of specified blocks, nse. This tensor encodes the index in values and row_indices depending on where the given row block starts. Each successive number in the tensor subtracted by the number before it denotes the number of blocks in a given column.

  • The row_indices tensor contains the row block indices of each element. This is a (B + 1)-D tensor of shape (*batchsize, nse).

  • The values tensor contains the values of the sparse BSC tensor elements collected into two-dimensional blocks. This is a (1 + 2 + K)-D tensor of shape (nse, nrowblocks, ncolblocks, *densesize).

BSC 张量的构造

可以使用 torch.sparse_bsc_tensor() 函数直接构造稀疏 BSC 张量。用户必须分别提供行和列块索引以及值张量，其中列块索引必须使用 CSR 压缩编码指定。如果 size 参数不存在，则会从 ccol_indices 和 row_indices 张量中推断出来。

>>> ccol_indices = torch.tensor([0, 2, 4])
>>> row_indices = torch.tensor([0, 1, 0, 1])
>>> values = torch.tensor([[[0, 1, 2], [6, 7, 8]],
...                        [[3, 4, 5], [9, 10, 11]],
...                        [[12, 13, 14], [18, 19, 20]],
...                        [[15, 16, 17], [21, 22, 23]]])
>>> bsc = torch.sparse_bsc_tensor(ccol_indices, row_indices, values, dtype=torch.float64)
>>> bsc
tensor(ccol_indices=tensor([0, 2, 4]),
       row_indices=tensor([0, 1, 0, 1]),
       values=tensor([[[ 0.,  1.,  2.],
                       [ 6.,  7.,  8.]],
                      [[ 3.,  4.,  5.],
                       [ 9., 10., 11.]],
                      [[12., 13., 14.],
                       [18., 19., 20.]],
                      [[15., 16., 17.],
                       [21., 22., 23.]]]), size=(4, 6), nnz=4,
       dtype=torch.float64, layout=torch.sparse_bsc)

处理稀疏压缩张量的工具

所有稀疏压缩张量——CSR、CSC、BSR 和 BSC 张量——在概念上非常相似，因为它们的索引数据被分成两部分：所谓的压缩索引使用 CSR 编码，而所谓的普通索引与压缩索引正交。这允许这些张量上的各种工具共享相同的实现，这些实现由张量布局参数化。

稀疏压缩张量的构造

稀疏 CSR、CSC、BSR 和 CSC 张量可以通过使用 torch.sparse_compressed_tensor() 函数来构造，该函数与上面讨论的构造函数 torch.sparse_csr_tensor()、torch.sparse_csc_tensor()、torch.sparse_bsr_tensor() 和 torch.sparse_bsc_tensor() 的接口相同，但有一个额外的必需的 layout 参数。以下示例说明了一种使用相同输入数据构造 CSR 和 CSC 张量的方法，方法是将相应的布局参数指定给 torch.sparse_compressed_tensor() 函数。

>>> compressed_indices = torch.tensor([0, 2, 4])
>>> plain_indices = torch.tensor([0, 1, 0, 1])
>>> values = torch.tensor([1, 2, 3, 4])
>>> csr = torch.sparse_compressed_tensor(compressed_indices, plain_indices, values, layout=torch.sparse_csr)
>>> 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,
       layout=torch.sparse_csr)
>>> csc = torch.sparse_compressed_tensor(compressed_indices, plain_indices, values, layout=torch.sparse_csc)
>>> csc
tensor(ccol_indices=tensor([0, 2, 4]),
       row_indices=tensor([0, 1, 0, 1]),
       values=tensor([1, 2, 3, 4]), size=(2, 2), nnz=4,
       layout=torch.sparse_csc)
>>> (csr.transpose(0, 1).to_dense() == csc.to_dense()).all()
tensor(True)

支持的操作

线性代数运算

下表总结了对稀疏矩阵支持的线性代数运算，其中操作数布局可能会有所不同。这里 T[layout] 表示具有给定布局的张量。类似地，M[layout] 表示矩阵（2 维 PyTorch 张量），V[layout] 表示向量（1 维 PyTorch 张量）。此外，f 表示标量（浮点数或 0 维 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.matmul()

no

M[SparseSemiStructured] @ M[strided] -> M[strided]

torch.matmul()

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.mm()

no

M[strided] @ M[SparseSemiStructured] -> 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]

其中“稀疏梯度？”列指示 PyTorch 操作是否支持相对于稀疏矩阵参数的向后传播。除 torch.smm() 外，所有 PyTorch 操作都支持相对于带步长矩阵参数的向后传播。

注意

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

张量方法和稀疏

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

Tensor.is_sparse

如果张量使用稀疏 COO 存储布局，则为 True，否则为 False。

Tensor.is_sparse_csr

如果张量使用稀疏 CSR 存储布局，则为 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

将张量转换为压缩行存储格式 (CSR)。

Tensor.to_sparse_csc

将张量转换为压缩列存储 (CSC) 格式。

Tensor.to_sparse_bsr

将张量转换为给定块大小的块稀疏行 (BSR) 存储格式。

Tensor.to_sparse_bsc

将张量转换为给定块大小的块稀疏列 (BSC) 存储格式。

Tensor.to_dense

如果 self 不是步幅张量，则创建 self 的步幅副本，否则返回 self。

Tensor.values

返回 稀疏 COO 张量 的值张量。

以下张量方法专门用于稀疏 COO 张量

Tensor.coalesce

如果 self 是 未合并张量，则返回 self 的合并副本。

Tensor.sparse_resize_

将 self 稀疏张量 调整为所需的尺寸以及稀疏和密集维数。

Tensor.sparse_resize_and_clear_

从 稀疏张量 self 中移除所有指定元素，并将 self 调整为所需的尺寸以及稀疏和密集维数。

Tensor.is_coalesced

如果 self 是一个已合并的 稀疏 COO 张量，则返回 True，否则返回 False。

Tensor.indices

返回 稀疏 COO 张量 的索引张量。

以下方法专门用于 稀疏 CSR 张量 和 稀疏 BSR 张量

Tensor.crow_indices

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

Tensor.col_indices

当 self 是布局为 sparse_csr 的稀疏 CSR 张量时，返回包含 self 张量的列索引的张量。

以下方法专门用于 稀疏 CSC 张量 和 稀疏 BSC 张量

Tensor.row_indices

Tensor.ccol_indices

以下张量方法支持稀疏 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_()

专门用于稀疏张量的 Torch 函数

sparse_coo_tensor

使用给定 indices 上的指定值构造 COO(rdinate) 格式的稀疏张量。

sparse_csr_tensor

使用给定 crow_indices 和 col_indices 上的指定值构造 CSR (压缩稀疏行) 格式的稀疏张量。

sparse_csc_tensor

使用给定 ccol_indices 和 row_indices 上的指定值构造 CSC (压缩稀疏列) 格式的稀疏张量。

sparse_bsr_tensor

使用指定的二维块在给定的 crow_indices 和 col_indices 处构造一个 BSR（块压缩稀疏行）格式的稀疏张量。

sparse_bsc_tensor

使用指定的二维块在给定的 ccol_indices 和 row_indices 处构造一个 BSC（块压缩稀疏列）格式的稀疏张量。

sparse_compressed_tensor

使用在给定的 compressed_indices 和 plain_indices 处指定的数值，构造一个 压缩稀疏格式 - CSR、CSC、BSR 或 BSC - 的稀疏张量。

sparse.sum

返回给定稀疏张量的每一行的总和。

sparse.addmm

此函数在正向传播中与 torch.addmm() 执行完全相同的功能，不同之处在于它支持稀疏 COO 矩阵 mat1 的反向传播。

sparse.sampled_addmm

在由 input 的稀疏模式指定的 位置执行稠密矩阵 mat1 和 mat2 的矩阵乘法。

sparse.mm

执行稀疏矩阵 mat1

sspaddmm

将稀疏张量 mat1 与稠密张量 mat2 进行矩阵乘法，然后将稀疏张量 input 添加到结果中。

hspmm

执行 稀疏 COO 矩阵 mat1 和带步幅矩阵 mat2 的矩阵乘法。

smm

执行稀疏矩阵 input 与稠密矩阵 mat 的矩阵乘法。

sparse.softmax

应用 softmax 函数。

sparse.log_softmax

应用 softmax 函数，然后取对数。

sparse.spdiags

通过将 diagonals 的行中的值放置在输出的指定对角线上，创建一个稀疏二维张量。

其他函数

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

要管理检查稀疏张量的不变性，请参阅

sparse.check_sparse_tensor_invariants

用于控制检查稀疏张量不变性的工具。

要将稀疏张量与 gradcheck() 函数一起使用，请参阅

sparse.as_sparse_gradcheck

装饰器函数，用于扩展稀疏张量的 gradcheck。

一元函数

我们的目标是支持所有保留零的一元函数。

如果您发现我们缺少您需要的保留零的一元函数，请随时提出功能请求，并打开一个问题。与往常一样，请在打开问题之前先尝试使用搜索功能。

以下运算符目前支持稀疏 COO/CSR/CSC/BSR/CSR 张量输入。

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

文档

访问 PyTorch 的综合开发者文档

查看文档

教程

获取面向初学者和高级开发者的深入教程

查看教程

资源

查找开发资源并获得问题的解答

查看资源

为了分析流量并优化您的体验，我们在本网站上提供 cookie。通过点击或导航，您同意允许我们使用 cookie。作为本网站的当前维护者，Facebook 的 cookie 政策适用。了解更多信息，包括可用的控制项：Cookie 政策。