torch.sparse¶
警告
PyTorch API 中的稀疏张量处于测试阶段，可能会在不久的将来发生变化。 我们非常欢迎功能请求、错误报告和一般建议作为 GitHub 问题提交。
为什么以及何时使用稀疏性¶
默认情况下，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 目前支持 COO, CSR, CSC, BSR, 和 BSC。
我们还提供了一个原型实现来支持 :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 "" , line 1, in 如上例所示，我们不支持非零保留的一元运算符，例如 cos。非零保留的一元运算的输出将无法像输入那样充分利用稀疏存储格式，并且可能会导致内存的灾难性增加。我们依赖用户先显式转换为密集张量，然后再运行该操作。
>>> b_s.to_dense().cos()
tensor([[ 1.0000, -0.4161],
[-0.9900, 1.0000]])
我们了解到一些用户希望在执行诸如cos等操作时忽略压缩的零，而不是保留操作的精确语义。为此，我们可以参考torch.masked及其MaskedTensor，后者同样由稀疏存储格式和内核支持。
另请注意，目前用户没有选择输出布局的选项。例如，将稀疏张量添加到常规的跨距张量会生成一个跨距张量。一些用户可能希望保持稀疏布局，因为他们知道结果仍然会足够稀疏。
>>> a + b.to_sparse()
张量([[0., 3.],
[3., 0.]])
我们承认，能够高效生成不同输出布局的内核访问是非常有用的。后续操作可能会显著受益于接收特定的布局。我们正在开发一个API来控制结果布局，并认识到这是一个重要的功能，可以为任何给定的模型规划更优化的执行路径。
稀疏半结构化张量¶
警告
稀疏半结构化张量目前是一个原型功能，可能会发生变化。请随时打开一个问题来报告错误或分享您的反馈。
半结构化稀疏性是一种稀疏数据布局，最早由NVIDIA在其Ampere架构中引入。它也被称为细粒度结构化稀疏性或2:4结构化稀疏性。
这种稀疏布局存储了每2n个元素中的n个元素，其中n由张量的数据类型（dtype）的宽度决定。最常用的dtype是float16，其中n=2，因此称为“2:4结构化稀疏性”。
半结构化稀疏性在这篇NVIDIA博客文章中有更详细的解释。
在 PyTorch 中，半结构化稀疏性是通过一个 Tensor 子类实现的。
通过子类化，我们可以覆盖 __torch_dispatch__，从而在执行矩阵乘法时使用更快的稀疏内核。
我们还可以在子类中以压缩形式存储张量，以减少内存开销。
在这种压缩形式中，稀疏张量通过仅保留指定的元素和一些元数据来存储，这些元数据编码了掩码。
注意
指定的元素和半结构化稀疏张量的元数据掩码一起存储在一个单一的扁平压缩张量中。它们被附加在一起，形成一个连续的内存块。
压缩张量 = [ 原始张量的指定元素 | 元数据掩码 ]
对于大小为 (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.shape 和 e = bitwidth(tensor.dtype)，因此对于 torch.float16 和 torch.bfloat16，e = 16，而对于 torch.int8，e = 8。
使用这些计算，我们可以确定原始密集表示和新稀疏表示的总内存占用。
这为我们提供了一个简单的压缩比公式，该公式仅依赖于张量数据类型的位宽。
通过使用这个公式，我们发现对于torch.float16或torch.bfloat16，压缩比为56.25%，而对于torch.int8，压缩比为62.5%。
构建稀疏半结构化张量¶
您可以通过使用 torch.to_sparse_semi_structured 函数将密集张量转换为稀疏半结构化张量。
请注意，由于半结构化稀疏性的硬件兼容性仅限于NVIDIA GPU，我们仅支持CUDA张量。
以下数据类型支持半结构化稀疏性。请注意，每种数据类型都有其自己的形状约束和压缩因子。
PyTorch 数据类型 |
形状约束 |
压缩因子 |
稀疏模式 |
|---|---|---|---|
|
张量必须是2D的，并且(r, c)都必须是64的正倍数 |
9/16 |
2:4 |
|
张量必须是2D的，并且(r, c)都必须是64的正倍数 |
9/16 |
2:4 |
|
张量必须是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()
张量([[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=张量([[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=张量([[-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(偏置, 密集, 稀疏.t())
torch.mm(dense, sparse)
torch.mm(稀疏, 密集)
aten.linear.default(dense, sparse, bias)
aten.t.default(稀疏)
aten.t.detach(稀疏)
要使用这些操作，只需传递 to_sparse_semi_structured(tensor) 的输出，而不是在张量以半结构化稀疏格式包含0时使用 tensor，如下所示：
>>> 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 + <元素类型的大小（以字节为单位）>) * nse 字节（加上存储其他张量数据的常数开销）。
一个跨步张量的内存消耗至少为
product(<张量 形状>) * <元素 类型 的字节大小>。
例如，一个10000 x 10000的张量，包含100000个非零的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))
>>> # 或者另一种等效的表达方式来获取 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。相应的（张量）值收集在
values大小为(nse, dense_dims)的张量中，并且具有任意整数 或浮点数元素类型。
注意
我们使用 (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()
张量([[[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
张量(indices=张量([[1, 1]]),
values=张量( [3, 4]),
size=(3,), nnz=2, layout=torch.sparse_coo)
在合并过程中，将使用求和将多值元素累积为单个值：
>>> s.coalesce()
张量(indices=张量([[1]]),
values=张量([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
张量(indices=张量([[0, 0, 1, 1]]),
values=张量([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()
运行时错误：无法获取未合并张量的索引，请先调用 .coalesce()
要获取未合并张量的COO格式数据，请使用
torch.Tensor._values() 和 torch.Tensor._indices()：
>>> s._indices()
张量([[0, 1, 1],
[2, 0, 2]])
警告
调用 torch.Tensor._values() 将返回一个 分离的 张量。
要跟踪梯度，必须使用 torch.Tensor.coalesce().values()。
构建一个新的稀疏COO张量会生成一个未合并的张量：
>>> s.is_coalesced()
假
但可以使用 torch.Tensor.coalesce() 方法构造稀疏 COO 张量的合并副本：
>>> s2 = s.coalesce()
>>> s2.indices()
张量([[0, 1, 1],
[2, 0, 2]])
当处理未合并的稀疏COO张量时，必须考虑未合并数据的加法性质：相同索引的值是求和的项，其求和结果给出相应张量元素的值。例如，稀疏未合并张量的标量乘法可以通过将所有未合并的值与标量相乘来实现，因为c *
(a + b) == c * a + c * b成立。然而，任何非线性操作，例如平方根，不能通过将操作应用于未合并的数据来实现，因为sqrt(a + b) == sqrt(a) + sqrt(b)通常不成立。
对稀疏COO张量的切片（正步长）仅支持密集维度。索引支持稀疏和密集维度：
>>> s[1]
张量(indices=张量([[0, 2]]),
values=张量([[5, 6],
[7, 8]]),
size=(3, 2), nnz=2, layout=torch.sparse_coo)
>>> s[1, 0, 1]
张量(6)
>>> s[1, 0, 1:]
张量([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个块被指定。该张量的values形状为(b, n, p, q)。这种对指定元素数量的解释来自于所有稀疏压缩布局都源自于二维矩阵的压缩。批次维度被视为稀疏矩阵的堆叠，密集维度将元素的含义从简单的标量值更改为具有自身维度的数组。
稀疏 CSR 张量¶
CSR 格式相对于 COO 格式的主要优势在于更好地利用存储空间和更快的计算操作，例如使用 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-D 张量。The
valuestensor 包含 CSR 张量的元素值。这是一个大小为nse的 1-D 张量。
注意
索引张量 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 + <元素类型的大小（以字节为单位）> *
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_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 张量构造 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_indices、row_indices和values：
ccol_indices张量由压缩列索引组成。这是一个形状为(*batchsize, ncols + 1)的 (B + 1)-D 张量。最后一个元素是指定的元素数量，nse。该张量根据给定列的起始位置对values和row_indices中的索引进行编码。张量中每个连续的数字减去前一个数字表示给定列中的元素数量。张量
row_indices包含每个元素的行索引。这是一个形状为(*batchsize, nse)的 (B + 1)-D 张量。The
valuestensor 包含 CSC tensor 元素的值。这是一个形状为(nse, *densesize)的 (1 + K)-D tensor。
CSC张量的构建¶
稀疏CSC张量可以直接通过使用
torch.sparse_csc_tensor() 函数来构建。用户必须分别提供行索引、列索引和值张量，其中列索引必须使用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 张量。(strided) 张量中的任何零将被解释为稀疏张量中的缺失值：
>>> 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)-D 张量。最后一个元素是指定的块数，nse。该张量根据给定列块的起始位置对values和col_indices中的索引进行编码。张量中每个连续的数字减去前一个数字表示给定行中的块数。张量
col_indices包含每个元素的列块索引。这是一个形状为(*batchsize, nse)的 (B + 1)-D 张量。张量
values包含了稀疏 BSR 张量元素的值，这些元素被收集到二维块中。这是一个形状为(nse, nrowblocks, ncolblocks, *densesize)的 (1 + 2 + K)-D 张量。
BSR张量的构建¶
稀疏BSR张量可以直接通过使用
torch.sparse_bsr_tensor() 函数来构建。用户必须分别提供行和列块索引以及值张量，其中行块索引必须使用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：
ccol_indices张量由压缩列索引组成。这是一个形状为(*batchsize, ncolblocks + 1)的 (B + 1)-D 张量。最后一个元素是指定的块数，nse。该张量根据给定行块的起始位置对values和row_indices中的索引进行编码。张量中每个连续的数字减去前一个数字表示给定列中的块数。张量
row_indices包含每个元素的行块索引。这是一个形状为(*batchsize, nse)的 (B + 1)-D 张量。The
valuestensor 包含稀疏 BSC 张量元素的值，这些值被收集到二维块中。这是一个形状为(nse, nrowblocks, ncolblocks, *densesize)的 (1 + 2 + K)-D 张量。
BSC张量的构建¶
稀疏BSC张量可以直接通过使用
torch.sparse_bsc_tensor() 函数来构建。用户必须分别提供行和列块索引以及值张量，其中列块索引必须使用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张量的方法：
>>> 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 操作 |
稀疏梯度？ |
布局签名 |
|---|---|---|
否 |
|
|
否 |
|
|
否 |
|
|
否 |
|
|
否 |
|
|
否 |
|
|
否 |
|
|
否 |
|
|
否 |
|
|
是 |
|
|
否 |
|
|
否 |
|
|
否 |
|
|
否 |
|
|
否 |
|
|
否 |
|
|
是 |
|
|
否 |
|
|
否 |
|
|
是 |
|
|
是 |
|
其中“稀疏梯度？”列指示 PyTorch 操作是否支持关于稀疏矩阵参数的反向传播。除 torch.smm() 外，所有 PyTorch 操作都支持关于密集矩阵参数的反向传播。
注意
目前，PyTorch 不支持使用布局签名 M[strided] @ M[sparse_coo] 进行矩阵乘法。然而，应用程序仍然可以使用矩阵关系 D @
S == (S.t() @ D.t()).t() 来计算。
张量方法和稀疏¶
以下Tensor方法与稀疏张量相关：
如果张量使用稀疏COO存储布局，则为 |
|
如果张量使用稀疏CSR存储布局，则为 |
|
返回稀疏张量中的稠密维度数量 |
|
返回一个稀疏张量中的稀疏维度数量 |
|
返回一个新的 稀疏张量，其值来自一个步幅张量 |
|
返回张量的稀疏副本。 |
|
将一个张量转换为坐标格式。 |
|
将张量转换为压缩行存储格式（CSR）。 |
|
将张量转换为压缩列存储（CSC）格式。 |
|
将张量转换为给定块大小的块稀疏行（BSR）存储格式。 |
|
将张量转换为给定块大小的块稀疏列（BSC）存储格式。 |
|
如果 |
|
返回一个稀疏COO张量的值张量。 |
以下 Tensor 方法是特定于稀疏 COO 张量的：
如果 |
|
调整 |
|
从稀疏张量 |
|
如果 |
|
返回一个稀疏COO张量的索引张量。 |
当 |
|
返回当 |
以下 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_()
特定于稀疏张量的Torch函数¶
sparse_coo_tensor |
构建一个COO(rdinate)格式的稀疏张量，并在给定的 |
sparse_csr_tensor |
构建一个CSR（压缩稀疏行）格式的稀疏张量，并在给定的 |
sparse_csc_tensor |
构建一个CSC（压缩稀疏列）格式的稀疏张量，并在给定的 |
sparse_bsr_tensor |
构建一个BSR（块压缩稀疏行）格式的稀疏张量，在给定的 |
sparse_bsc_tensor |
构建一个BSC（块压缩稀疏列）格式的稀疏张量，并在给定的 |
sparse_compressed_tensor |
构建一个压缩稀疏格式 - CSR、CSC、BSR 或 BSC - 的稀疏张量，并在给定的 |
返回给定稀疏张量每行的和。 |
|
此函数在前向传播中与 |
|
在由 |
|
对稀疏矩阵 |
|
sspaddmm |
矩阵将稀疏张量 |
hspmm |
对一个稀疏COO矩阵 |
smm |
对稀疏矩阵 |
应用一个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()
要管理检查稀疏张量不变量，请参见：
一个用于控制检查稀疏张量不变性的工具。 |
要使用稀疏张量与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()