torch.sparse#
创建时间: 2017年4月26日 | 最后更新时间: 2025年6月18日
警告
PyTorch 的稀疏张量 API 仍处于 Beta 阶段,未来可能会发生更改。我们非常欢迎将新功能请求、错误报告和一般性建议作为 GitHub issues 提交。
为什么要以及何时使用稀疏性#
默认情况下,PyTorch 将 torch.Tensor
的元素连续存储在物理内存中。这使得各种需要快速访问元素的数组处理算法能够得到高效的实现。
现在,一些用户可能会决定用大部分元素值为零的张量来表示图邻接矩阵、剪枝权重或点云等数据。我们认识到这些是重要的应用场景,并致力于通过稀疏存储格式为这些用例提供性能优化。
多年来,人们开发了各种稀疏存储格式,例如 COO、CSR/CSC、半结构化、LIL 等。虽然它们在具体布局上有所不同,但都通过有效地表示零值元素来压缩数据。我们称未压缩的值为指定的元素,与未指定的、被压缩的元素相对。
通过压缩重复的零值,稀疏存储格式旨在节省 CPU 和 GPU 上的内存和计算资源。特别是在稀疏度很高或稀疏性结构高度的情况下,这可能会对性能产生显著影响。因此,稀疏存储格式可以被视为一种性能优化。
与其他许多性能优化一样,稀疏存储格式并非总是有利的。当您尝试将稀疏格式用于您的用例时,您可能会发现执行时间反而增加而不是减少。
如果您在理论上预期性能会有显著提升,但实际测量到的性能却有所下降,请随时在 GitHub 上提交 issue。这有助于我们优先实现高效的内核和更广泛的性能优化。
我们使尝试不同的稀疏布局以及在它们之间进行转换变得容易,同时不会对您的特定应用程序的最佳选择发表意见。
功能概览#
我们希望通过为每种布局提供转换例程,使从给定的密集张量构建稀疏张量变得简单直接。
在下面的示例中,我们将一个默认密集(跨步)布局的 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: 半结构化稀疏<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 来控制结果布局,并认识到它是为任何给定模型规划更优执行路径的重要功能。
稀疏半结构化张量#
警告
稀疏半结构化张量目前是一项原型功能,可能会发生更改。如果您发现错误或有任何反馈,请随时提交 issue。
半结构化稀疏是一种稀疏数据布局,最初是在 NVIDIA 的 Ampere 架构中引入的。它也被称为**细粒度结构化稀疏**或**2:4 结构化稀疏**。
此稀疏布局存储每 2n 个元素中的 n 个元素,其中 n 由张量数据类型 (dtype) 的宽度决定。最常用的 dtype 是 float16,其中 n=2,因此称为“2:4 结构化稀疏”。
NVIDIA 的这篇博客文章 更详细地解释了半结构化稀疏。
在 PyTorch 中,半结构化稀疏通过张量子类实现。通过子类化,我们可以重写 `__torch_dispatch__`,允许我们在执行矩阵乘法时使用更快的稀疏内核。我们还可以将张量存储在其压缩形式中,以减少内存开销。
在此压缩形式中,稀疏张量仅通过保留*指定的*元素和一些元数据来存储,这些元数据编码了掩码。
注意
半结构化稀疏张量的指定元素和元数据掩码存储在单个扁平的压缩张量中。它们相互追加以形成连续的内存块。
压缩张量 = [原始张量的指定元素 | 元数据掩码]
对于大小为 (r, c) 的原始张量,我们期望前 m * k // 2 个元素是保留的元素,其余部分是元数据。
为了方便用户查看指定元素和掩码,可以使用 `.indices()` 和 `.values()` 分别访问掩码和指定元素。
`.values()` 以大小为 (r, c//2) 的张量形式返回指定元素,其 dtype 与密集矩阵相同。
`.indices()` 以大小为 (r, c//2) 的张量形式返回元数据掩码,如果 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。
使用这些计算,我们可以确定原始密集和新稀疏表示的总内存占用。
这为我们提供了一个简单的压缩比公式,它仅取决于张量数据类型的位宽。
使用此公式,我们发现 `torch.float16` 或 `torch.bfloat16` 的压缩比为 56.25%,`torch.int8` 的压缩比为 62.5%。
构造稀疏半结构化张量#
您可以通过简单地使用 `torch.to_sparse_semi_structured` 函数将密集张量转换为稀疏半结构化张量。
另请注意,我们仅支持 CUDA 张量,因为半结构化稀疏的硬件兼容性仅限于 NVIDIA GPU。
半结构化稀疏支持以下数据类型。请注意,每种数据类型都有其自身的形状约束和压缩因子。
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()
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`,如下所示:
>>> 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.int64` 的 `indices` 张量中。
相应的*值*收集在大小为 `(nse,)` 且元素类型为任意整数或浮点数的 `values` 张量中。
其中 `ndim` 是张量的维度,`nse` 是指定元素的数量。
注意
稀疏 COO 张量的内存消耗至少为 `(ndim * 8 +
跨步张量的内存消耗至少为 `product(
例如,一个包含 100,000 个非零 32 位浮点数的 10,000 x 10,000 张量,在使用 COO 张量布局时,内存消耗至少为 `(2 * 8 + 4) * 100,000 = 2,000,000` 字节;而在使用默认跨步张量布局时,内存消耗为 `10,000 * 10,000 * 4 = 400,000,000` 字节。注意,使用 COO 存储格式可节省 200 倍的内存。
构造#
可以通过向函数 torch.sparse_coo_tensor()
提供索引和值的两个张量,以及稀疏张量的大小(当无法从索引和值张量推断时)来构造稀疏 COO 张量。
假设我们要定义一个稀疏张量,其中条目 3 在位置 (0, 2),条目 4 在位置 (1, 0),条目 5 在位置 (1, 2)。未指定的元素假定具有相同的值(填充值),默认情况下为零。然后我们将这样写:
>>> 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 张量扩展了稀疏 COO 张量,允许 `values` 张量成为多维张量,因此我们有:
指定元素的索引收集在大小为 `(sparse_dims, nse)` 且元素类型为 `torch.int64` 的 `indices` 张量中。
相应的(张量)值收集在大小为 `(nse, dense_dims)` 且元素类型为任意整数或浮点数的 `values` 张量中。
注意
我们使用 (M + K) 维张量来表示 N 维稀疏混合张量,其中 M 和 K 分别是稀疏维度和密集维度的数量,且 M + K == N。
假设我们要创建一个 (2 + 1) 维张量,其中条目 [3, 4] 在位置 (0, 2),条目 [5, 6] 在位置 (1, 0),条目 [7, 8] 在位置 (1, 2)。我们将这样写:
>>> 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`,这会导致一个一维未合并张量。
>>> 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。稀疏压缩张量的稀疏维度数量始终为 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 张量#
CSR 格式相对于 COO 格式的主要优点是更好的存储利用率和更快的数据运算,例如使用 MKL 和 MAGMA 后端进行的稀疏矩阵-向量乘法。
在最简单的情况下,一个 (0 + 2 + 0) 维稀疏 CSR 张量由三个一维张量组成:`crow_indices`、`col_indices` 和 `values`。
`crow_indices` 张量包含压缩的行索引。这是一个大小为 `nrows + 1`(行数加 1)的一维张量。`crow_indices` 的最后一个元素是指定元素的数量 `nse`。此张量根据给定行在 `values` 和 `col_indices` 中的索引位置进行编码,指示给定行的起始位置。张量中连续的两个数字之差表示给定行中的元素数量。
`col_indices` 张量包含每个元素的列索引。这是一个大小为 `nse` 的一维张量。
`values` 张量包含 CSR 张量元素的值。这是一个大小为 `nse` 的一维张量。
注意
索引张量 `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 +
使用 稀疏 COO 格式介绍中的说明 中的相同示例数据,一个包含 100,000 个非零 32 位浮点数的 10,000 x 10,000 张量,在使用 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 张量构造 2D 稀疏 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` 张量包含压缩的列索引。这是一个 (B + 1) 维张量,形状为 `(*batchsize, ncols + 1)`。最后一个元素是指定元素的数量 `nse`。此张量根据给定列在 `values` 和 `row_indices` 中的索引位置进行编码,指示给定列的起始位置。张量中连续的两个数字之差表示给定列中的元素数量。
`row_indices` 张量包含每个元素的行索引。这是一个 (B + 1) 维张量,形状为 `(*batchsize, nse)`。
`values` 张量包含 CSC 张量元素的值。这是一个 (1 + K) 维张量,形状为 `(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 张量构造函数在行索引参数之前具有压缩列索引参数。
(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_indices`、`col_indices` 和 `values`。
`crow_indices` 张量包含压缩的行块索引。这是一个 (B + 1) 维张量,形状为 `(*batchsize, nrowblocks + 1)`。最后一个元素是指定块的数量 `nse`。此张量根据给定列块在 `values` 和 `col_indices` 中的索引位置进行编码,指示给定行块的起始位置。张量中连续的两个数字之差表示给定行中的块数量。
`col_indices` 张量包含每个元素的列块索引。这是一个 (B + 1) 维张量,形状为 `(*batchsize, nse)`。
`values` 张量包含收集到二维块中的稀疏 BSR 张量元素的*值*。这是一个 (1 + 2 + K) 维张量,形状为 `(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)
(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_indices`、`row_indices` 和 `values`。
`ccol_indices` 张量包含压缩的列块索引。这是一个 (B + 1) 维张量,形状为 `(*batchsize, ncolblocks + 1)`。最后一个元素是指定块的数量 `nse`。此张量根据给定行块在 `values` 和 `row_indices` 中的索引位置进行编码,指示给定列块的起始位置。张量中连续的两个数字之差表示给定列中的块数量。
`row_indices` 张量包含每个元素的行块索引。这是一个 (B + 1) 维张量,形状为 `(*batchsize, nse)`。
`values` 张量包含收集到二维块中的稀疏 BSC 张量元素的*值*。这是一个 (1 + 2 + K) 维张量,形状为 `(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 编码;以及所谓的普通索引,与压缩索引正交。这使得这些张量的各种工具可以共享相同的实现,并通过张量布局进行参数化。
稀疏压缩张量构造#
稀疏 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]` 表示矩阵(2D PyTorch 张量),`V[layout]` 表示向量(1D PyTorch 张量)。此外,`f` 表示标量(浮点数或 0D PyTorch 张量),`*` 表示逐元素乘法,`@` 表示矩阵乘法。
PyTorch 运算 |
稀疏梯度? |
布局签名 |
---|---|---|
否 |
|
|
否 |
|
|
否 |
|
|
否 |
|
|
否 |
|
|
否 |
|
|
否 |
|
|
否 |
|
|
否 |
|
|
是 |
|
|
否 |
|
|
否 |
|
|
否 |
|
|
否 |
|
|
否 |
|
|
否 |
|
|
是 |
|
|
否 |
|
|
否 |
|
|
否 |
|
|
是 |
|
|
是 |
|
其中“稀疏梯度?”列表示 PyTorch 运算是否支持相对于稀疏矩阵的后向传播。所有 PyTorch 运算(除了 `torch.smm()`)都支持相对于跨步矩阵的后向传播。
注意
目前,PyTorch 不支持使用布局签名 M[strided] @ M[sparse_coo]
的矩阵乘法。然而,应用程序仍然可以使用矩阵关系 D @ S == (S.t() @ D.t()).t()
来计算它。
Tensor 方法和稀疏#
以下 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 函数,特定于稀疏张量#
在指定的 |
|
在指定的 |
|
在指定的 |
|
在指定的 |
|
在指定的 |
|
使用指定的 |
|
返回给定稀疏张量每一行的和。 |
|
此函数在正向传播中执行的操作与 |
|
在 |
|
执行稀疏矩阵 |
|
将稀疏张量 |
|
执行 稀疏 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。 |
保持零的一元函数#
我们的目标是支持所有“保持零的一元函数”:将零映射到零的单参数函数。
如果您发现我们缺少您需要的一个保持零的一元函数,请随时为功能请求打开一个 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()