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()
