torch.sparse ¬
警告
PyTorch 的稀疏张量 API 处于测试阶段,未来可能会发生变化。我们非常欢迎通过 GitHub 问题报告功能提出功能请求、错误报告和一般建议。
为什么以及何时使用稀疏性 ¬
默认情况下,PyTorch 将 torch.Tensor
元素连续存储在物理内存中。这导致各种需要快速访问元素的数组处理算法实现高效。
现在,一些用户可能会选择用元素大部分为零的 Tensor 来表示图邻接矩阵、剪枝权重或点云等数据。我们认识到这些是重要的应用,并旨在通过稀疏存储格式对这些用例提供性能优化。
近年来,已经开发了各种稀疏存储格式,如 COO、CSR/CSC、半结构化、LIL 等。虽然它们的布局各不相同,但它们都通过高效地表示零值元素来压缩数据。我们称与未指定、压缩元素相对的未压缩值为“未压缩值”。
通过压缩重复的零值,稀疏存储格式旨在在各种 CPU 和 GPU 上节省内存和计算资源。特别是对于高稀疏度或高度结构化稀疏度,这可以产生显著的性能影响。因此,稀疏存储格式可以被视为一种性能优化。
许多其他性能优化稀疏存储格式并不总是有优势。在尝试为您的用例使用稀疏格式时,您可能会发现执行时间增加而不是减少。
如果您在分析上预期会看到性能显著提升,但测量结果却是下降,请鼓励您在 GitHub 上创建一个 issue。这有助于我们优先实施高效的内核和更广泛性能优化。
我们使您能够轻松尝试不同的稀疏布局,并在它们之间进行转换,而不会对您特定应用的最佳选择有所偏见。
功能概述 ¶
我们希望能够通过提供每个布局的转换例程,使从给定的密集 Tensor 构建稀疏 Tensor 变得简单直接。
在下一个示例中,我们将默认的密集(带偏移量)布局的 2D Tensor 转换为基于 COO 内存布局的 2D Tensor。在这种情况下,仅存储非零元素的值和索引。
>>> 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。请参阅参考资料以获取更多详细信息。
注意,我们对这些格式进行了一些轻微的概括。
批处理:如 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 步进 Tensor 创建一个具有 2 个稀疏维度和 1 个密集维度的 3D 混合 COO 张量。如果 3D 步进 Tensor 中的整行都是零,则不会存储。然而,如果行中的任何值非零,则整个行都会被存储。这减少了索引的数量,因为我们只需要为每行提供一个索引,而不是为每个元素提供一个索引。但这也会增加值的存储量。由于只有整行都是零的行才能被发出,任何非零值的存在都会导致整行被存储。
>>> 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)
运算符概述
基本上,对稀疏存储格式的 Tensor 的操作与对步进(或其它)存储格式的 Tensor 的操作行为相同。存储的特定性,即数据的物理布局,会影响操作的性能,但不应该影响语义。
我们正在积极增加稀疏张量的运算符覆盖范围。用户不应期望获得与密集 Tensor 相同的支持水平。请参阅我们的运算符文档以获取列表。
>>> 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。非零保留的算术运算结果无法像输入那样充分利用稀疏存储格式,可能会造成内存的灾难性增长。我们反而依赖用户显式地将 Tensor 转换为稠密格式,然后再执行操作。
>>> b_s.to_dense().cos()
tensor([[ 1.0000, -0.4161],
[-0.9900, 1.0000]])
我们知道,一些用户希望忽略压缩零,例如在 cos 运算中而不是保留操作的精确语义。为此,我们可以指向 torch.masked 及其 MaskedTensor,它反过来也由稀疏存储格式和内核支持。
还要注意,目前用户无法选择输出布局。例如,将稀疏 Tensor 与常规的 strided Tensor 相加,结果将是 strided Tensor。一些用户可能希望保持稀疏布局,因为他们知道结果仍然足够稀疏。
>>> 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__
,这样我们可以在执行矩阵乘法时使用更快的稀疏内核。我们还可以在子类中将张量以压缩形式存储,以减少内存开销。
在这种压缩形式中,稀疏张量通过仅保留指定的元素和一些元数据(编码掩码)来存储。
注意
半结构化稀疏张量的指定元素和元数据掩码存储在一个单独的扁平压缩张量中。它们被依次连接,形成一个连续的内存块。
压缩张量 = [原始张量的指定元素 | 元数据掩码]
对于大小为 (r, c) 的原始张量,我们期望前 m * k // 2 个元素是保留元素,其余的张量是元数据。
为了使用户更容易查看指定的元素和掩码,可以使用 .indices()
和 .values()
分别访问掩码和指定元素。
返回一个大小为 (r, c//2) 的张量,其数据类型与密集矩阵相同。
对于 2:4 稀疏张量,元数据开销很小 - 每个指定元素仅占用 2 位。
对于 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 数据类型 |
形状约束 |
压缩因子 |
稀疏模式 |
---|---|---|---|
|
张量必须是二维的,且(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(稠密矩阵, 稀疏矩阵)
torch.mm(稀疏矩阵, 稠密矩阵)
aten.linear.default(稠密矩阵, 稀疏矩阵, 偏置)
aten.t.default(稀疏矩阵)
aten.t.detach(sparse)
要使用这些操作,只需传递 to_sparse_semi_structured(tensor)
的输出,而不是使用 tensor
,一旦你的张量以半结构化稀疏格式包含 0,如下所示:
>>> a = torch.Tensor([0, 0, 1, 1]).tile((64, 16)).half().cuda()
>>> b = torch.rand(64, 64).half().cuda()
>>> c = torch.mm(a, b)
>>> a_sparse = to_sparse_semi_structured(a)
>>> torch.allclose(c, torch.mm(a_sparse, b))
True
加速具有半结构化稀疏性的 nn.Linear ¶
如果权重已经是半结构化稀疏的,你可以用几行代码加速模型中的线性层:
>>> input = torch.rand(64, 64).half().cuda()
>>> mask = torch.Tensor([0, 0, 1, 1]).tile((64, 16)).cuda().bool()
>>> linear = nn.Linear(64, 64).half().cuda()
>>> linear.weight = nn.Parameter(to_sparse_semi_structured(linear.weight.masked_fill(~mask, 0)))
稀疏 COO 张量 ¶
PyTorch 实现了所谓的坐标格式,或称 COO 格式,作为实现稀疏张量的存储格式之一。在 COO 格式中,指定的元素存储为元素索引和对应值的元组。特别是,
指定元素的索引收集在大小为
(ndim, nse)
的indices
张量中,元素类型为torch.int64
,对应的值收集在大小为
(nse,)
的values
张量中,元素类型为任意整数或浮点数,
其中 ndim
是张量的维度, nse
是指定的元素数量。
注意
稀疏 COO 张量的内存消耗至少为 (ndim *
8 + <size of element type in bytes>) * nse
字节(加上存储其他张量数据产生的常数开销)。
稀疏张量的内存消耗至少为 product(<tensor shape>) * <size of element type in bytes>
。
例如,使用 COO 张量布局时,一个包含 100,000 个非零 32 位浮点数的 10,000 x 10,000 张量的内存消耗至少为 (2 * 8 + 4) * 100 000 = 2 000 000
字节,而使用默认的步进张量布局时至少为 10 000 * 10 000 * 4 = 400 000 000
字节。请注意,使用 COO 存储格式可以节省 200 倍的内存。
构建
可以通过提供索引张量和值张量以及稀疏张量的大小(当无法从索引和值张量中推断出来时)到一个函数 torch.sparse_coo_tensor()
来构建一个稀疏的 COO 张量。
假设我们想定义一个稀疏张量,其中包含位置(0,2)处的条目 3,位置(1,0)处的条目 4,以及位置(1,2)处的条目 5。未指定的元素默认具有相同的值,填充值,默认为零。然后我们可以这样写:
>>> i = [[0, 1, 1],
[2, 0, 2]]
>>> v = [3, 4, 5]
>>> s = torch.sparse_coo_tensor(i, v, (2, 3))
>>> s
tensor(indices=tensor([[0, 1, 1],
[2, 0, 2]]),
values=tensor([3, 4, 5]),
size=(2, 3), nnz=3, layout=torch.sparse_coo)
>>> s.to_dense()
tensor([[0, 0, 3],
[4, 0, 5]])
注意,输入 i
不是一个索引元组的列表。如果你想以这种方式写入你的索引,你应该在传递给稀疏构造函数之前进行转置:
>>> i = [[0, 2], [1, 0], [1, 2]]
>>> v = [3, 4, 5 ]
>>> s = torch.sparse_coo_tensor(list(zip(*i)), v, (2, 3))
>>> # Or another equivalent formulation to get s
>>> s = torch.sparse_coo_tensor(torch.tensor(i).t(), v, (2, 3))
>>> torch.sparse_coo_tensor(i.t(), v, torch.Size([2,3])).to_dense()
tensor([[0, 0, 3],
[4, 0, 5]])
可以仅指定其大小来构造一个空的稀疏 COO 张量:
>>> torch.sparse_coo_tensor(size=(2, 3))
tensor(indices=tensor([], size=(2, 0)),
values=tensor([], size=(0,)),
size=(2, 3), nnz=0, layout=torch.sparse_coo)
稀疏混合 COO 张量
PyTorch 实现了稀疏张量向稀疏张量(连续)张量的扩展,具有标量值的稀疏张量称为混合张量。
PyTorch 混合 COO 张量通过允许 values
张量成为一个多维张量来扩展稀疏 COO 张量,因此我们有:
指定元素的索引被收集在大小为
(sparse_dims, nse)
、元素类型为torch.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
的张量维度是稀疏和密集维度数量的总和,- 稀疏索引是显式存储的,
- 混合张量的值是 K 维张量,
- 值以步长张量的形式存储。
注意
稠密维度始终跟在稀疏维度之后,即不支持稠密和稀疏维度的混合。
注意
要确保构建的稀疏张量具有一致的索引、值和大小,可以在创建张量时通过 check_invariants=True
关键字参数启用不变性检查,或者使用 torch.sparse.check_sparse_tensor_invariants
上下文管理器实例全局启用。默认情况下,稀疏张量的不变性检查是禁用的。
未合并的稀疏 COO 张量
PyTorch 稀疏 COO 张量格式允许存在未合并的稀疏张量,其中索引中可能存在重复的坐标;在这种情况下,解释是,该索引处的值是所有重复值条目的总和。例如,可以为相同的索引 1
指定多个值,从而得到一个 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 张量,则其 COO 格式数据可以通过方法 torch.Tensor.indices()
和 torch.Tensor.values()
获取。
注意
目前,只有在张量实例合并时才能获取 COO 格式数据:
>>> s.indices()
RuntimeError: Cannot get indices on an uncoalesced tensor, please call .coalesce() first
要获取未合并张量的 COO 格式数据,请使用 torch.Tensor._values()
和 torch.Tensor._indices()
:
>>> s._indices()
tensor([[0, 1, 1],
[2, 0, 2]])
警告
调用 torch.Tensor._values()
将返回一个分离的张量。要跟踪梯度,必须使用 torch.Tensor.coalesce().values()
。
构建一个新的稀疏 COO 张量会产生一个未合并的张量:
>>> s.is_coalesced()
False
但可以使用 torch.Tensor.coalesce()
方法构建稀疏 COO 张量的合并副本:
>>> s2 = s.coalesce()
>>> s2.indices()
tensor([[0, 1, 1],
[2, 0, 2]])
当处理未合并的稀疏 COO 张量时,必须考虑到未合并数据的加性:相同索引的值是求和的项,求和的结果给出了相应张量元素的值。例如,对稀疏未合并张量进行标量乘法可以通过将所有未合并的值乘以标量来实现,因为 c *
(a + b) == c * a + c * b
成立。然而,任何非线性操作,例如平方根,不能通过将操作应用于未合并数据来实现,因为在一般情况下 sqrt(a + b) == sqrt(a) + sqrt(b)
不成立。
稀疏 COO 张量的切片(正步长)仅支持密集维度。索引支持稀疏和密集维度:
>>> s[1]
tensor(indices=tensor([[0, 2]]),
values=tensor([[5, 6],
[7, 8]]),
size=(3, 2), nnz=2, layout=torch.sparse_coo)
>>> s[1, 0, 1]
tensor(6)
>>> s[1, 0, 1:]
tensor([6])
在 PyTorch 中,稀疏张量的填充值不能显式指定,通常默认为零。然而,存在一些操作可能会以不同的方式解释填充值。例如, torch.sparse.softmax()
计算 softmax 时假定填充值为负无穷。
稀疏压缩张量
稀疏压缩张量代表一类稀疏张量,它们具有一个共同特征,即使用一种编码来压缩特定维度的索引,这种编码使得对稀疏压缩张量的线性代数核进行某些优化成为可能。这种编码基于压缩稀疏行(CSR)格式,PyTorch 稀疏压缩张量扩展了这种格式,支持稀疏张量批处理,允许多维张量值,并将稀疏张量值存储在密集块中。
注意
我们使用(B + M + K)维张量来表示一个 N 维稀疏压缩混合张量,其中 B、M 和 K 分别是批处理、稀疏和密集维度的数量,使得 B + M + K == N
成立。稀疏压缩张量的稀疏维度数量始终为两个, M == 2
。
注意
我们说,如果一个索引张量 compressed_indices
满足以下不变性,则它使用 CSR 编码:
compressed_indices
是一个连续的步长为 32 或 64 位的整数张量
compressed_indices
的形状是(*batchsize, compressed_dim_size + 1)
,其中compressed_dim_size
表示压缩的维度数量(例如行或列)compressed_indices[..., 0] == 0
,其中...
表示批处理索引compressed_indices[..., compressed_dim_size] == nse
表示指定元素的数量
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 维块被视为指定的元素。以一个三维块稀疏张量为例,它有一个长度为 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
:
crow_indices
张量包含压缩的行索引。这是一个大小为nrows + 1
(行数加 1)的 1-D 张量。crow_indices
的最后一个元素是指定元素的数量,nse
。这个张量根据给定行的起始位置编码values
和col_indices
中的索引。张量中连续的数字减去前面的数字表示给定行中的元素数量。
col_indices
张量包含每个元素的列索引。这是一个大小为nse
的 1-D 张量。
values
索引张量包含 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 + <size of element type in bytes> *
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 张量的构建
可以通过使用 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 张量的 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 张量。
values
索引张量包含 CSC 张量元素的值。这是一个形状为(nse, *densesize)
的(1 + K)-D 张量。
CSC 张量的构建
可以直接使用 torch.sparse_csc_tensor()
函数构建稀疏的 CSC 张量。用户必须分别提供行索引和列索引以及值张量,其中列索引必须使用 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()
方法从任何二维张量构建。在(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)
(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)-D 张量。最后一个元素是指定块的数目,nse
。这个张量根据给定的行块开始位置编码索引values
和row_indices
。张量中连续的数字减去前面的数字表示给定列中的块数。
row_indices
张量包含每个元素的行块索引。这是一个形状为(*batchsize, nse)
的 (B + 1)-D 张量。
values
张量包含稀疏 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_csr_tensor()
、 torch.sparse_csc_tensor()
、 torch.sparse_bsr_tensor()
和 torch.sparse_bsc_tensor()
具有相同接口的 torch.sparse_compressed_tensor()
函数构建,但需要额外的 layout
参数。以下示例说明了使用相同输入数据构建 CSR 和 CSC 张量的方法,通过指定相应的布局参数给 torch.sparse_compressed_tensor()
函数:
>>> 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()
来计算。
张量方法和稀疏
以下张量方法与稀疏张量相关:
如果张量使用稀疏 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_()
火炬函数针对稀疏张量特定功能
在指定的 |
|
在指定的 |
|
在指定的 |
|
构建一个在指定 |
|
构建一个在指定 |
|
在 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()