torch.sparse#
创建日期:2017 年 4 月 26 日 | 最后更新日期:2025 年 6 月 18 日
警告
稀疏张量的 PyTorch API 目前处于测试阶段,未来可能会有所变动。我们非常欢迎通过 GitHub Issues 提交功能需求、错误报告和通用建议。
为何使用以及何时使用稀疏性#
默认情况下,PyTorch 将 torch.Tensor 元素在物理内存中连续存储。这使得各种需要快速访问元素的数组处理算法能够得到高效的实现。
然而,一些用户可能会决定用张量表示诸如图邻接矩阵、剪枝权重或点云等数据,这些张量中 绝大部分元素的值为零。我们认识到这些是重要的应用场景,并旨在通过稀疏存储格式为这些用例提供性能优化。
多年来,人们开发了多种稀疏存储格式,如 COO、CSR/CSC、半结构化(semi-structured)、LIL 等。虽然它们在具体布局上有所不同,但它们都通过对零值元素的高效表示来压缩数据。我们将这些压缩后的元素称为 非指定(unspecified),与之相对的则是 指定(specified) 的非压缩值。
通过压缩重复的零值,稀疏存储格式旨在节省各类 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>。请参阅参考资料以获取更多详细信息。
请注意,我们对这些格式进行了轻微的泛化处理。
批处理(Batching):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)
稠密维度(Dense dimensions):另一方面,某些数据(例如图嵌入)可能更适合视为向量的稀疏集合,而不是标量的集合。
在此示例中,我们从一个 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 架构中引入的稀疏数据布局。它也被称为 细粒度结构化稀疏性(fine-grained structured sparsity) 或 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)
要使用这些算子,一旦你的张量在半结构化稀疏格式中包含 0,只需传递 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 将所谓的坐标格式(Coordinate format,即 COO 格式)作为实现稀疏张量的存储格式之一。在 COO 格式中,指定元素以元素索引元组及其对应值的方式存储。具体而言:
指定元素的索引被收集在大小为
(ndim, nse)、元素类型为torch.int64的indices张量中;相应的值被收集在大小为
(nse,)、元素类型为任意整数或浮点数的values张量中;
其中 ndim 是张量的维度,nse 是指定元素的数量。
注意
稀疏 COO 张量的内存消耗至少为 (ndim * 8 + <元素类型的字节大小>) * nse 字节(加上存储其他张量数据的常量开销)。
步长张量的内存消耗至少为 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 倍内存节省。
构造#
稀疏 COO 张量可以通过向 torch.sparse_coo_tensor() 函数提供索引和值这两个张量,以及稀疏张量的大小(当无法从索引和值张量中推断时)来构建。
假设我们想要定义一个稀疏张量,其中位置 (0, 2) 处条目为 3,位置 (1, 0) 处条目为 4,位置 (1, 2) 处条目为 5。未指定元素被假定具有相同的填充值(fill value),默认情况下为零。我们可以这样写:
>>> 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.int64的indices张量中;相应的(张量)值收集在大小为
(nse, dense_dims)、元素类型为任意整数或浮点数的values张量中。
注意
我们使用 (M + K) 维张量来表示 N 维稀疏混合张量,其中 M 和 K 分别是稀疏维度和稠密维度的数量,满足 M + K == N。
假设我们要创建一个 (2 + 1) 维张量,其中位置 (0, 2) 处条目为 [3, 4],位置 (1, 0) 处条目为 [5, 6],位置 (1, 2) 处条目为 [7, 8]。我们可以这样写:
>>> i = [[0, 1, 1],
[2, 0, 2]]
>>> v = [[3, 4], [5, 6], [7, 8]]
>>> s = torch.sparse_coo_tensor(i, v, (2, 3, 2))
>>> s
tensor(indices=tensor([[0, 1, 1],
[2, 0, 2]]),
values=tensor([[3, 4],
[5, 6],
[7, 8]]),
size=(2, 3, 2), nnz=3, layout=torch.sparse_coo)
>>> s.to_dense()
tensor([[[0, 0],
[0, 0],
[3, 4]],
[[5, 6],
[0, 0],
[7, 8]]])
通常,如果 s 是一个稀疏 COO 张量,且 M = s.sparse_dim(),K = s.dense_dim(),则存在以下不变式:
M + K == len(s.shape) == s.ndim—— 张量的维度是稀疏维度和稠密维度数量之和;
s.indices().shape == (M, nse)—— 稀疏索引被显式存储;
s.values().shape == (nse,) + s.shape[M : M + K]—— 混合张量的值是 K 维张量;
s.values().layout == torch.strided—— 值作为步长张量存储。
注意
稠密维度总是跟在稀疏维度之后,也就是说,不支持稠密维度和稀疏维度的混合。
注意
为确保构造的稀疏张量具有一致的索引、值和大小,可以通过 check_invariants=True 关键字参数在创建张量时启用不变式检查,或使用 torch.sparse.check_sparse_tensor_invariants 上下文管理器实例全局启用。默认情况下,稀疏张量不变式检查是禁用的。
未合并(Uncoalesced)的稀疏 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()),则应适时对稀疏张量进行合并(coalesce),以防止它们变得过大。
另一方面,索引的字典序对于实现涉及许多元素选择操作(如切片或矩阵乘法)的算法可能具有优势。
使用稀疏 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 张量,那么它的 COO 格式数据可以通过方法 torch.Tensor.indices() 和 torch.Tensor.values() 获取。
注意
目前,只有当张量实例已合并(coalesced)时,才能获取其 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() 将返回一个脱离计算图(detached)的张量。若要跟踪梯度,必须改用 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 中,稀疏张量的填充值(fill value)不能显式指定,通常默认为零。但是,有些操作对填充值的解释可能不同。例如,torch.sparse.softmax() 计算 softmax 时,假设填充值为负无穷大。
稀疏压缩张量#
稀疏压缩张量是一类稀疏张量,其共同特点是使用特定编码压缩某一维度的索引,从而能够针对稀疏压缩张量的线性代数核进行优化。这种编码基于 压缩稀疏行 (CSR) 格式,PyTorch 稀疏压缩张量在此基础上扩展了对稀疏张量批次的支持,允许存在多维张量值,并将稀疏张量值存储在稠密块中。
注意
我们使用 (B + M + K) 维张量来表示 N 维稀疏压缩混合张量,其中 B、M 和 K 分别是批次维度、稀疏维度和稠密维度的数量,且满足 B + M + K == N。稀疏压缩张量的稀疏维度数量始终为二,即 M == 2。
注意
如果索引张量 compressed_indices 满足以下不变量,我们称其使用 CSR 编码
compressed_indices是一个连续的跨步(strided)32 位或 64 位整数张量。compressed_indices的形状为(*batchsize, compressed_dim_size + 1),其中compressed_dim_size是压缩维度的数量(例如行数或列数)。compressed_indices[..., 0] == 0,其中...表示批次索引。compressed_indices[..., compressed_dim_size] == nse,其中nse是指定元素的数量。对于
i=1, ..., compressed_dim_size,满足0 <= compressed_indices[..., i] - compressed_indices[..., i - 1] <= plain_dim_size,其中plain_dim_size是普通维度(与压缩维度正交,例如列或行)的数量。
为确保构造的稀疏张量具有一致的索引、值和大小,可以通过 check_invariants=True 关键字参数在创建张量时启用不变式检查,或使用 torch.sparse.check_sparse_tensor_invariants 上下文管理器实例全局启用。默认情况下,稀疏张量不变式检查是禁用的。
注意
稀疏压缩布局到 N 维张量的推广可能会导致对指定元素数量的混淆。当稀疏压缩张量包含批次维度时,指定元素的数量对应于每个批次中的元素数量。当稀疏压缩张量具有稠密维度时,被考虑的元素现在是 K 维数组。此外,对于分块稀疏压缩布局,二维块被视为被指定的元素。以一个 3 维块稀疏张量为例,它具有一个长度为 b 的批次维度,以及一个形状为 p, q 的块形状。如果该张量有 n 个指定元素,那么实际上我们每个批次指定了 n 个块。该张量的 values 形状将为 (b, n, p, q)。这种对指定元素数量的解释源于所有稀疏压缩布局均衍生自二维矩阵的压缩。批次维度被视为稀疏矩阵的堆叠,稠密维度将元素的含义从简单的标量值改变为具有自身维度的数组。
稀疏 CSR 张量#
与 COO 格式相比,CSR 格式的主要优势在于更好地利用了存储空间,并且在 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 默认链接的是使用 32 位整数索引的 MKL LP64。
在一般情况下,(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 格式介绍中的注释 里的相同示例数据,当使用 CSR 张量布局时,一个 10 000 x 10 000、包含 100 000 个 32 位浮点非零元素的张量,其内存消耗至少为 (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 张量构建二维稀疏 CSR 张量的最简单方法是使用 torch.Tensor.to_sparse_csr() 方法。(跨步)张量中的任何零都将被解释为稀疏张量中的缺失值。
>>> a = torch.tensor([[0, 0, 1, 0], [1, 2, 0, 0], [0, 0, 0, 0]], dtype=torch.float64)
>>> sp = a.to_sparse_csr()
>>> sp
tensor(crow_indices=tensor([0, 1, 3, 3]),
col_indices=tensor([2, 0, 1]),
values=tensor([1., 1., 2.]), size=(3, 4), nnz=3, dtype=torch.float64)
CSR 张量运算#
稀疏矩阵向量乘法可以使用 tensor.matmul() 方法执行。这是目前 CSR 张量上支持的唯一数学运算。
>>> vec = torch.randn(4, 1, dtype=torch.float64)
>>> sp.matmul(vec)
tensor([[0.9078],
[1.3180],
[0.0000]], dtype=torch.float64)
稀疏 CSC 张量#
稀疏 CSC(压缩稀疏列)张量格式实现了用于存储二维张量的 CSC 格式,并扩展支持了稀疏 CSC 张量批次,且值可以是多维张量。
注意
当转置涉及交换稀疏维度时,稀疏 CSC 张量本质上是稀疏 CSR 张量的转置。
与 稀疏 CSR 张量 类似,稀疏 CSC 张量由三个张量组成:ccol_indices、row_indices 和 values。
ccol_indices张量由压缩列索引组成。这是一个形状为(*batchsize, ncols + 1)的 (B + 1) 维张量。最后一个元素是指定元素的数量nse。该张量根据给定列的起始位置编码了values和row_indices中的索引。张量中每个连续数字减去其前一个数字,即表示给定列中的元素数量。
row_indices张量包含每个元素的行索引。这是一个形状为(*batchsize, nse)的 (B + 1) 维张量。
values张量包含 CSC 张量元素的值。这是一个形状为(nse, *densesize)的 (1 + K) 维张量。
CSC 张量的构建#
稀疏 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张量由压缩行索引组成。这是一个形状为(*batchsize, nrowblocks + 1)的 (B + 1) 维张量。最后一个元素是指定块的数量nse。该张量根据给定列块的起始位置编码了values和col_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_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张量由压缩列索引组成。这是一个形状为(*batchsize, ncolblocks + 1)的 (B + 1) 维张量。最后一个元素是指定块的数量nse。该张量根据给定行块的起始位置编码了values和row_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_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 和 BSC 张量可以使用 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] 表示矩阵(二维 PyTorch 张量),V[layout] 表示向量(一维 PyTorch 张量)。此外,f 表示标量(浮点数或 0 维 PyTorch 张量),* 为逐元素乘法,@ 为矩阵乘法。
PyTorch 运算 |
支持稀疏梯度? |
布局签名 |
|---|---|---|
否 |
|
|
否 |
|
|
否 |
|
|
否 |
|
|
否 |
|
|
否 |
|
|
否 |
|
|
否 |
|
|
否 |
|
|
是 |
|
|
否 |
|
|
否 |
|
|
否 |
|
|
否 |
|
|
否 |
|
|
否 |
|
|
是 |
|
|
否 |
|
|
否 |
|
|
否 |
|
|
是 |
|
|
是 |
|
其中“Sparse grad?”列表示该 PyTorch 操作是否支持针对稀疏矩阵参数的梯度反向传播。除了 torch.smm() 外,所有 PyTorch 操作均支持针对步长(strided)矩阵参数的梯度反向传播。
注意
目前,PyTorch 不支持布局签名为 M[strided] @ M[sparse_coo] 的矩阵乘法。但是,应用程序仍可通过矩阵关系 D @ S == (S.t() @ D.t()).t() 来进行计算。
Tensor 方法与稀疏张量#
以下 Tensor 方法与稀疏张量相关
如果张量使用稀疏 COO 存储布局,则为 |
|
如果 Tensor 使用稀疏 CSR 存储布局,则为 |
|
返回 稀疏张量 |
|
返回 稀疏张量 |
|
返回一个新的 稀疏张量,其值来自步长张量 |
|
返回该张量的稀疏副本。 |
|
将张量转换为 坐标格式(COO)。 |
|
将张量转换为压缩稀疏行(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_()
专门用于稀疏张量的 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。 |
保持零值的单变量函数#
我们的目标是支持所有“保持零值的单变量函数”:即接受一个参数并将零映射为零的函数。
如果您发现我们缺少您所需的此类函数,请随时发起功能请求。一如既往,请在发起请求前先尝试使用搜索功能。
以下运算符当前支持稀疏 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()