torch.sparse¶
警告
PyTorch 稀疏张量的 API 处于测试阶段，未来可能会发生变化。我们非常欢迎通过 GitHub 问题提交功能请求、错误报告和一般建议。
何时以及为何使用稀疏性¶
默认情况下，PyTorch会在物理内存中连续存储torch.Tensor个元素。这使得各种需要快速访问元素的数组处理算法能够高效实现。
现在，一些用户可能决定使用张量来表示图邻接矩阵、剪枝后的权重或点云数据，这些张量的元素大部分为零。我们认识到这些是重要的应用场景，并希望通过稀疏存储格式为这些用例提供性能优化。
多年来，已经开发了各种稀疏存储格式，例如 COO、CSR/CSC、半结构化、LIL 等。尽管它们在具体布局上有所不同，但它们都通过高效表示零值元素来压缩数据。我们将未压缩的值称为已指定，与未指定的压缩元素相对。
通过压缩重复的零值，稀疏存储格式旨在节省各种CPU和GPU上的内存和计算资源。特别是在高稀疏度或高度结构化稀疏的情况下，这可能会产生显著的性能影响。因此，稀疏存储格式可以被视为一种性能优化。
像许多其他性能优化的稀疏存储格式一样，并不总是有利的。在尝试为你的用例使用稀疏格式时，你可能会发现执行时间增加而不是减少。
如果您在分析后预期性能会有显著提升，但实际测量结果却出现了下降，请务必打开一个 GitHub 问题。这有助于我们优先实现高效的内核以及更广泛的性能优化。
我们使尝试不同的稀疏布局变得简单，并且可以在它们之间进行转换，而不会对您特定应用程序的最佳选择发表意见。
功能概述¶
我们希望从给定的密集张量构造稀疏张量变得简单，为此我们将为每种布局提供转换例程。
在下一个示例中，我们将一个默认的二维密集（strided）布局张量转换为由 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: semi-structured sparsity<sparse-semi-structured-docs>. 请参阅参考资料以获取更多详细信息。
请注意，我们提供了这些格式的轻微泛化版本。
批处理：像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 中，半结构化稀疏性通过 Tensor 子类实现。
通过子类化，我们可以覆盖 __torch_dispatch__ ，从而在执行矩阵乘法时使用更快的稀疏内核。
我们还可以在子类中以压缩形式存储张量，以减少内存开销。
在这种压缩形式中，稀疏张量通过仅保留指定的元素和一些元数据进行存储，这些元数据对掩码进行了编码。
注意
指定的元素和元数据掩码在一个半结构化稀疏张量中存储在一起，作为一个单一的扁平压缩张量。它们被附加在一起形成一个连续的内存块。
压缩张量 = [ 原始张量的指定元素 | 元数据_mask ]
对于大小为 (r, c) 的原始张量，我们期望保留前 m * k // 2 个元素， 其余的张量是元数据。
为了使用户更容易查看指定的元素和掩码，可以使用 .indices() 和 .values() 分别访问掩码和指定元素。
.values()返回一个大小为 (r, c//2) 的张量，其中包含指定的元素，并且与密集矩阵具有相同的 dtype。.indices()如果 dtype 是 torch.float16 或 torch.bfloat16，则返回一个大小为 (r, c//2 ) 且元素类型为torch.int16的 tensor 中的 metadata_mask；如果 dtype 是 torch.int8，则元素类型为torch.int32。
对于 2:4 稀疏张量，元数据开销很小 - 每个指定元素仅需 2 位。
注意
需要注意的是，torch.float32 仅支持 1:2 稀疏性。因此，它不遵循与上述相同的公式。
我们在这里分解如何计算 2:4 稀疏张量的压缩率（密集尺寸 / 稀疏尺寸）。
设 (r, c) = tensor.shape 和 e = bitwidth(tensor.dtype)，因此 e = 16 对于 torch.float16 和 torch.bfloat16 以及 e = 8 对于 torch.int8。
通过这些计算，我们可以确定原始密集表示和新稀疏表示的总内存占用情况。
这为我们提供了一个简单的压缩率公式，该公式仅取决于张量数据类型的位宽。
通过使用这个公式，我们发现压缩比对于torch.float16或torch.bfloat16是56.25%，而对于torch.int8是62.5%。
构建稀疏半结构化张量¶
你可以通过简单地使用torch.to_sparse_semi_structured函数，将密集张量转换为稀疏半结构化张量。
请注意，由于半结构化稀疏性的硬件兼容性仅限于NVIDIA GPU，因此我们仅支持CUDA张量。
以下数据类型支持半结构化稀疏性。请注意，每种数据类型都有其自身的形状约束和压缩系数。
PyTorch 数据类型 |
形状约束 |
压缩因子 |
稀疏模式 |
|---|---|---|---|
|
张量必须是二维的，且 (r, c) 必须都是 64 的正整数倍 |
9/16 |
2:4 |
|
张量必须是二维的，且 (r, c) 必须都是 64 的正整数倍 |
9/16 |
2:4 |
|
张量必须是二维的，且 (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(稀疏, 密集)
aten.linear.default(dense, sparse, bias)
aten.t.default(sparse)
aten.t.detach(sparse)
要使用这些操作，只需传递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 格式中，指定的元素以元素索引和对应值的元组形式存储。特别地，
the indices of specified elements are collected in
indicestensor of size(ndim, nse)and with element typetorch.int64,the corresponding values are collected in
valuestensor 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
indicestensor of size(sparse_dims, nse)and with element typetorch.int64,the corresponding (tensor) values are collected in
valuestensor 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.
注意
密集维度总是跟在稀疏维度之后，也就是说，不支持密集和稀疏维度的混合。
注意
为了确保构建的稀疏张量具有一致的索引、值和大小，可以通过 check_invariants=True 关键字参数在每个张量创建时启用不变性检查，或者使用 torch.sparse.check_sparse_tensor_invariants 上下文管理器实例全局启用。默认情况下，稀疏张量的不变性检查是禁用的。
未合并的稀疏 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）格式，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_sizefori=1, ..., compressed_dim_size, whereplain_dim_sizeis the number of plain dimensions (orthogonal to compressed dimensions, e.g. columns or rows).
为了确保构建的稀疏张量具有一致的索引、值和大小，可以通过 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-D 张量组成：crow_indices、col_indices 和
values：
The
crow_indicestensor consists of compressed row indices. This is a 1-D tensor of sizenrows + 1(the number of rows plus 1). The last element ofcrow_indicesis the number of specified elements,nse. This tensor encodes the index invaluesandcol_indicesdepending 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_indicestensor contains the column indices of each element. This is a 1-D tensor of sizense.The
valuestensor contains the values of the CSR tensor elements. This is a 1-D tensor of sizense.
注意
索引张量 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 格式介绍中的笔记 相同的示例数据，一个包含 100,000 个非零 32 位浮点数的 10,000 x 10,000 张量在使用 CSR 张量布局时，内存消耗至少为 (10000 * 8 + (8 + 4 * 1) * 100 000) * 1 = 1 280 000 字节。请注意，与使用 COO 和 strided 格式相比，使用 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：
The
ccol_indicestensor consists of compressed column indices. This is a (B + 1)-D tensor of shape(*batchsize, ncols + 1). The last element is the number of specified elements,nse. This tensor encodes the index invaluesandrow_indicesdepending on where the given column starts. Each successive number in the tensor subtracted by the number before it denotes the number of elements in a given column.The
row_indicestensor contains the row indices of each element. This is a (B + 1)-D tensor of shape(*batchsize, nse).The
valuestensor contains the values of the CSC tensor elements. This is a (1 + K)-D tensor of shape(nse, *densesize).
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 张量。在（分块）张量中的任何零值都将被解释为稀疏张量中的缺失值：
>>> 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：
The
crow_indicestensor consists of compressed row indices. This is a (B + 1)-D tensor of shape(*batchsize, nrowblocks + 1). The last element is the number of specified blocks,nse. This tensor encodes the index invaluesandcol_indicesdepending on where the given column block starts. Each successive number in the tensor subtracted by the number before it denotes the number of blocks in a given row.The
col_indicestensor contains the column block indices of each element. This is a (B + 1)-D tensor of shape(*batchsize, nse).The
valuestensor 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 张量¶
稀疏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：
The
ccol_indicestensor 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 invaluesandrow_indicesdepending 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_indicestensor contains the row block indices of each element. This is a (B + 1)-D tensor of shape(*batchsize, nse).The
valuestensor 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 张量的构建¶
稀疏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 编码，而所谓的普通索引与压缩索引正交。这使得这些张量上的各种工具可以共享由张量布局参数化的相同实现。
稀疏压缩张量的构建¶
Sparse 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 操作 |
稀疏梯度？ |
布局签名 |
|---|---|---|
no |
|
|
no |
|
|
no |
|
|
no |
|
|
no |
|
|
no |
|
|
no |
|
|
no |
|
|
no |
|
|
是的 |
|
|
no |
|
|
no |
|
|
no |
|
|
no |
|
|
no |
|
|
no |
|
|
是的 |
|
|
no |
|
|
no |
|
|
no |
|
|
是的 |
|
|
是的 |
|
其中“是否稀疏梯度？”列表示 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 张量 的索引张量。 |
以下方法特定于 稀疏的 CSR 张量 和 稀疏的 BSR 张量：
返回包含 |
|
返回包含 |
以下方法特定于 稀疏的 CSC 张量 和 稀疏的 BSC 张量：
以下 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函数¶
构造一个 以COO（坐标）格式表示的稀疏张量，并在给定的 |
|
构造一个 在 CSR（压缩稀疏行）格式下的稀疏张量，并在给定的 |
|
构造一个 在 CSC（压缩稀疏列）格式下的稀疏张量，并在给定的 |
|
构造一个 在 BSR（Block Compressed Sparse Row）格式下的稀疏张量，并在给定的 |
|
构造一个在 BSC（块压缩稀疏列）格式下的稀疏张量，指定在给定的 |
|
构造一个在压缩稀疏格式（CSR、CSC、BSR 或 BSC）中具有指定值的 稀疏张量，这些值位于给定的 |
|
返回给定稀疏张量的每一行的总和。 |
|
此函数在前向传播中的作用与 |
|
执行密集矩阵 |
|
对稀疏矩阵 |
|
矩阵将稀疏张量 |
|
执行稀疏 COO 矩阵 |
|
对稀疏矩阵 |
|
应用一个 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()