评价此页

torch.sparse#

创建日期:2017年4月26日 | 最后更新日期:2025年6月18日

警告

PyTorch 稀疏张量 API 处于测试阶段,近期可能会有更改。我们非常欢迎通过 GitHub Issues 提出功能请求、错误报告和一般建议。

为何以及何时使用稀疏性#

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

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

多年来,已经开发了各种稀疏存储格式,如 COO、CSR/CSC、半结构化、LIL 等。尽管它们的精确布局不同,但它们都通过零值的有效表示来压缩数据。我们将未压缩的值称为**指定值**,以区别于**未指定值**(压缩元素)。

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

与许多其他性能优化一样,稀疏存储格式并非总是具有优势。当尝试将稀疏格式用于您的用例时,您可能会发现执行时间反而增加了。

如果您分析上预期会看到性能显著提升,但实际测量却出现下降,请随时在 GitHub 上提出问题。这将帮助我们优先实现高效的内核和更广泛的性能优化。

我们让尝试不同的稀疏布局并进行转换变得容易,而不对您的特定应用程序的最佳选择持固有偏见。

功能概述#

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

在下一个示例中,我们将具有默认稠密(步进)布局的 2D 张量转换为由 COO 内存布局支持的 2D 张量。在这种情况下,仅存储非零元素的值和索引。

>>> 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 目前支持 COOCSRCSCBSRBSC

我们还有一个支持 :ref:半结构化稀疏性 的原型实现。请参阅参考资料了解更多详情。

请注意,我们对这些格式进行了一些泛化。

批处理:GPU 等设备需要批处理才能获得最佳性能,因此我们支持批处理维度。

我们目前提供了一个非常简单的批处理版本,其中稀疏格式的每个组件本身都被批处理。这也要求每个批处理条目具有相同数量的指定元素。在此示例中,我们从 3D 稠密张量构造 3D(批处理)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)

稠密维度:另一方面,一些数据(如图嵌入)可能更适合看作稀疏向量集合而不是标量。

在此示例中,我们从 3D 步进张量创建了一个具有 2 个稀疏维度和 1 个稠密维度的 3D 混合 COO 张量。如果 3D 步进张量中的整行都为零,则不存储。但是,如果行中的任何值非零,则会完全存储。这减少了索引的数量,因为每行只需要一个索引,而不是每个元素一个。但这也会增加值的存储量。因为只有**完全**为零的行才能被省略,任何非零值的存在都会导致整行被存储。

>>> 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__,从而在使用矩阵乘法时使用更快的稀疏内核。我们还可以在子类中以压缩形式存储张量,以减少内存开销。

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

注意

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

压缩张量 = [原始张量的指定元素 | metadata_mask]

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

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

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

  • .indices() 返回大小为 (r, c//2 ) 且元素类型为 torch.int16(如果 dtype 为 torch.float16 或 torch.bfloat16)或 torch.int32(如果 dtype 为 torch.int8)的元数据掩码。

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

注意

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

在此,我们详细说明如何计算 2:4 稀疏张量的压缩比(稠密大小/稀疏大小)。

(r, c) = tensor.shapee = bitwidth(tensor.dtype),因此对于 torch.float16torch.bfloat16e = 16;对于 torch.int8e = 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.float16torch.bfloat16,压缩比为 56.25%,对于 torch.int8,压缩比为 62.5%。

构建稀疏半结构化张量#

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

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

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

PyTorch dtype

形状约束

压缩因子

稀疏模式

torch.float16

张量必须是二维的,并且 (r, c) 都必须是 64 的正倍数

9/16

2:4

torch.bfloat16

张量必须是二维的,并且 (r, c) 都必须是 64 的正倍数

9/16

2:4

torch.int8

张量必须是二维的,并且 (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) 且元素类型为 torch.int64indices 张量中,

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

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

注意

稀疏 COO 张量的内存消耗至少为 (ndim * 8 + <元素类型大小(字节)>) * nse 字节(加上存储其他张量数据的固定开销)。

步幅张量的内存消耗至少为 product(<张量形状>) * <元素类型大小(字节)>

例如,一个 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 张量,因此我们有

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

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

注意

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

假设我们要在位置 (0, 2) 创建一个值为 [3, 4],在位置 (1, 0) 创建一个值为 [5, 6],在位置 (1, 2) 创建一个值为 [7, 8] 的 (2 + 1) 维张量。我们可以这样写

>>> 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 指定多个值 34,从而形成一个 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() 方法的输出是一个具有以下属性的稀疏张量

注意

在大多数情况下,您不需要关心稀疏张量是否合并,因为大多数操作在合并或未合并的稀疏张量上都能相同地工作。

然而,有些操作可以在未合并张量上更有效地实现,有些则在合并张量上更有效地实现。

例如,稀疏 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_sparsetorch.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 张量,那么它的 COO 格式数据可以使用方法 torch.Tensor.indices()torch.Tensor.values() 获取。

注意

目前,只有当张量实例已合并时,才能获取 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 成立。稀疏压缩张量的稀疏维度数量总是两个,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 维数组。此外,对于块稀疏压缩布局,2-D 块被视为正在指定的元素。例如,一个 3 维块稀疏张量,其中一个批处理维度的长度为 b,块形状为 p, q。如果此张量有 n 个指定元素,那么实际上我们每个批处理指定了 n 个块。此张量将具有形状为 (b, n, p, q)values。这种对指定元素数量的解释源于所有稀疏压缩布局都源自 2 维矩阵的压缩。批处理维度被视为稀疏矩阵的堆叠,密集维度将元素的含义从一个简单的标量值更改为具有其自身维度的数组。

稀疏 CSR 张量#

CSR 格式相对于 COO 格式的主要优势在于更好地利用存储空间,以及使用 MKL 和 MAGMA 后端进行稀疏矩阵-向量乘法等计算操作速度更快。

在最简单的情况下,一个 (0 + 2 + 0) 维稀疏 CSR 张量由三个 1-D 张量组成:crow_indicescol_indicesvalues

  • crow_indices 张量包含压缩行索引。这是一个大小为 nrows + 1(行数加 1)的 1-D 张量。crow_indices 的最后一个元素是指定元素的数量,nse。此张量根据给定行开始的位置对 valuescol_indices 中的索引进行编码。张量中每个连续数字减去其之前的数字表示给定行中的元素数量。

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

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

注意

索引张量 crow_indicescol_indices 的元素类型应为 torch.int64(默认)或 torch.int32。如果要使用启用 MKL 的矩阵操作,请使用 torch.int32。这是因为 pytorch 的默认链接是 MKL LP64,它使用 32 位整数索引。

在一般情况下,(B + 2 + K) 维稀疏 CSR 张量由两个 (B + 1) 维索引张量 crow_indicescol_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) == Blen(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 + <元素类型字节大小> * prod(densesize)) * nse) * prod(batchsize) 字节(加上存储其他张量数据的固定开销)。

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

CSR 张量的构造#

稀疏 CSR 张量可以直接使用 torch.sparse_csr_tensor() 函数构造。用户必须分别提供行和列索引以及值张量,其中行索引必须使用 CSR 压缩编码指定。size 参数是可选的,如果不存在,将从 crow_indicescol_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 张量构造 2-D 稀疏 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_indicesrow_indicesvalues

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

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

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

CSC 张量的构造#

稀疏 CSC 张量可以直接使用 torch.sparse_csc_tensor() 函数构造。用户必须分别提供行和列索引以及值张量,其中列索引必须使用 CSR 压缩编码指定。size 参数是可选的,如果不存在,将从 row_indicesccol_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 张量构造函数将压缩列索引参数置于行索引参数之前。

(0 + 2 + 0) 维稀疏 CSC 张量可以使用 torch.Tensor.to_sparse_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_indicescol_indicesvalues

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

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

  • values 张量包含收集到二维块中的稀疏 BSR 张量元素的值。这是一个形状为 (nse, nrowblocks, ncolblocks, *densesize) 的 (1 + 2 + K) 维张量。

BSR 张量的构造#

稀疏 BSR 张量可以直接使用 torch.sparse_bsr_tensor() 函数构造。用户必须分别提供行和列块索引以及值张量,其中行块索引必须使用 CSR 压缩编码指定。size 参数是可选的,如果不存在,将从 crow_indicescol_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)

(0 + 2 + 0) 维稀疏 BSR 张量可以使用 torch.Tensor.to_sparse_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_indicesrow_indicesvalues

  • ccol_indices 张量包含压缩列索引。这是一个形状为 (*batchsize, ncolblocks + 1) 的 (B + 1) 维张量。最后一个元素是指定块的数量,nse。此张量根据给定行块的起始位置在 valuesrow_indices 中编码索引。张量中每个连续数字减去其之前的数字表示给定列中的块数量。

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

  • values 张量包含收集到二维块中的稀疏 BSC 张量元素的值。这是一个形状为 (nse, nrowblocks, ncolblocks, *densesize) 的 (1 + 2 + K) 维张量。

BSC 张量的构造#

稀疏 BSC 张量可以直接使用 torch.sparse_bsc_tensor() 函数构造。用户必须分别提供行和列块索引以及值张量,其中列块索引必须使用 CSR 压缩编码指定。size 参数是可选的,如果不存在,将从 ccol_indicesrow_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 参数。以下示例说明了使用相同输入数据通过指定相应的布局参数给 torch.sparse_compressed_tensor() 函数来构造 CSR 和 CSC 张量的方法

>>> 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-D PyTorch 张量),V[layout] 表示向量(1-D PyTorch 张量)。此外,f 表示标量(浮点数或 0-D PyTorch 张量),* 是逐元素乘法,@ 是矩阵乘法。

PyTorch 操作

稀疏梯度?

布局签名

torch.mv()

M[sparse_coo] @ V[strided] -> V[strided]

torch.mv()

M[sparse_csr] @ V[strided] -> V[strided]

torch.matmul()

M[sparse_coo] @ M[strided] -> M[strided]

torch.matmul()

M[sparse_csr] @ M[strided] -> M[strided]

torch.matmul()

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

torch.matmul()

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

torch.mm()

M[sparse_coo] @ M[strided] -> M[strided]

torch.mm()

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

torch.mm()

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

torch.sparse.mm()

M[sparse_coo] @ M[strided] -> M[strided]

torch.smm()

M[sparse_coo] @ M[strided] -> M[sparse_coo]

torch.hspmm()

M[sparse_coo] @ M[strided] -> M[hybrid sparse_coo]

torch.bmm()

T[sparse_coo] @ T[strided] -> T[strided]

torch.addmm()

f * M[strided] + f * (M[sparse_coo] @ M[strided]) -> M[strided]

torch.addmm()

f * M[strided] + f * (M[SparseSemiStructured] @ M[strided]) -> M[strided]

torch.addmm()

f * M[strided] + f * (M[strided] @ M[SparseSemiStructured]) -> M[strided]

torch.sparse.addmm()

f * M[strided] + f * (M[sparse_coo] @ M[strided]) -> M[strided]

torch.sparse.spsolve()

SOLVE(M[sparse_csr], V[strided]) -> V[strided]

torch.sspaddmm()

f * M[sparse_coo] + f * (M[sparse_coo] @ M[strided]) -> M[sparse_coo]

torch.lobpcg()

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.is_sparse

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

如果张量使用稀疏 CSR 存储布局,则为 True,否则为 False

Tensor.is_sparse_csr

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 张量

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

以 COO(坐标) 格式构建稀疏张量,并在给定 indices 处指定值。

sparse_csr_tensor

以 CSR (压缩稀疏行) 格式构建稀疏张量,并在给定 crow_indicescol_indices 处指定值。

sparse_csc_tensor

以 CSC (压缩稀疏列) 格式构建稀疏张量,并在给定 ccol_indicesrow_indices 处指定值。

sparse_bsr_tensor

以 BSR (块压缩稀疏行) 格式构建稀疏张量,并在给定 crow_indicescol_indices 处指定二维块。

sparse_bsc_tensor

以 BSC (块压缩稀疏列) 格式构建稀疏张量,并在给定 ccol_indicesrow_indices 处指定二维块。

sparse_compressed_tensor

以压缩稀疏格式(CSR、CSC、BSR 或 BSC)构建稀疏张量,并在给定 compressed_indicesplain_indices 处指定值。

sparse.sum

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

sparse.addmm

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

sparse.sampled_addmm

在由 input 的稀疏模式指定的位置执行密集矩阵 mat1mat2 的矩阵乘法。

sparse.mm

执行稀疏矩阵 mat1 的矩阵乘法。

sspaddmm

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

hspmm

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

smm

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

sparse.softmax

应用 softmax 函数。

sparse.spsolve

计算具有唯一解的方线性方程组的解。

sparse.log_softmax

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

sparse.spdiags

通过将 diagonals 行中的值沿着输出的指定对角线放置,创建稀疏 2D 张量。

其他函数#

以下 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。

零保留一元函数#

我们的目标是支持所有“零保留一元函数”:即那些将零映射到零的单参数函数。

如果您发现我们缺少您需要的零保留一元函数,请随时打开一个功能请求 issue。一如既往,请在打开 issue 之前先尝试搜索功能。

以下运算符目前支持稀疏 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()