torch.sparse¶
警告
稀疏張量的 PyTorch API 處於 Beta 階段,並且在不久的將來可能會發生變化。我們非常歡迎將功能請求、錯誤報告和一般建議作為 GitHub issue 提出。
為何以及何時使用稀疏性¶
預設情況下,PyTorch 會將 torch.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: 半結構化稀疏性<sparse-semi-structured-docs>。請參閱參考資料以取得更多詳細資訊。
請注意,我們提供了這些格式的輕微概括。
批次處理:諸如 GPU 之類的裝置需要批次處理才能獲得最佳效能,因此我們支援批次維度。
我們目前提供一個非常簡單的批次處理版本,其中稀疏格式的每個組件本身都是批次的。這也需要每個批次條目具有相同數量的指定元素。在此範例中,我們從 3D 密集 Tensor 構造一個 3D(批次處理)CSR Tensor。
>>> 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 Tensor。如果 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 的運算子覆蓋率。使用者不應期望與密集 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 新增到常規跨步 Tensor 會產生跨步 Tensor。有些使用者可能希望它保持稀疏佈局,因為他們知道結果仍然會足夠稀疏。
>>> a + b.to_sparse()
tensor([[0., 3.],
[3., 0.]])
我們承認,存取可以有效產生不同輸出佈局的核心非常有用。後續運算可能會從接收特定佈局中受益匪淺。我們正在開發一個 API 來控制結果佈局,並且認識到它是一個重要的功能,可以為任何給定的模型規劃更佳的執行路徑。
稀疏半結構化 Tensor¶
警告
稀疏半結構化 Tensor 目前是一個原型功能,可能會發生變化。請隨時開啟一個 issue 來報告錯誤或如果您有回饋要分享。
半結構化稀疏性是一種稀疏資料佈局,最初在 NVIDIA 的 Ampere 架構中引入。它也稱為細粒度結構化稀疏性或 2:4 結構化稀疏性。
此稀疏佈局儲存每 2n 個元素中的 n 個元素,其中 n 由 Tensor 的資料類型 (dtype) 的寬度決定。最常用的 dtype 是 float16,其中 n=2,因此稱為「2:4 結構化稀疏性」。
有關半結構化稀疏性的更詳細說明,請參閱 此 NVIDIA 部落格文章。
在 PyTorch 中,半結構化稀疏性是透過 Tensor 子類別實作的。透過子類別化,我們可以覆寫 __torch_dispatch__
,讓我們可以在執行矩陣乘法時使用更快的稀疏核心。我們也可以將 Tensor 以壓縮形式儲存在子類別中,以減少記憶體開銷。
在這個壓縮形式中,稀疏 Tensor 僅保留指定的元素和一些元資料(編碼遮罩)來儲存。
注意
半結構化稀疏 Tensor 的指定元素和元資料遮罩一起儲存在單個扁平壓縮 Tensor 中。它們彼此附加以形成連續的記憶體區塊。
壓縮的 tensor = [ 原始 tensor 的指定元素 | metadata_mask ]
對於大小為 (r, c) 的原始 tensor,我們預期前 m * k // 2 個元素是要保留的元素,而 tensor 的其餘部分是元資料。
為了讓使用者更容易檢視指定元素和遮罩,可以使用 .indices()
和 .values()
分別存取遮罩和指定元素。
.values()
傳回大小為 (r, c//2) 且與密集矩陣具有相同 dtype 的 tensor 中的指定元素。.indices()
傳回大小為 (r, c//2 ) 且元素類型為torch.int16
(如果 dtype 為 torch.float16 或 torch.bfloat16)以及元素類型為torch.int32
(如果 dtype 為 torch.int8)的 metadata_mask。
對於 2:4 稀疏 Tensor,元資料開銷很小 - 每個指定元素只有 2 個位元。
注意
請注意,torch.float32
僅支援 1:2 稀疏性。因此,它不適用於上述公式。
在這裡,我們將分解如何計算 2:4 稀疏張量的壓縮率(密集大小 / 稀疏大小)。
令 (r, c) = tensor.shape 且 e = bitwidth(tensor.dtype),因此 e = 16 對於 torch.float16
和 torch.bfloat16
,以及 e = 8 對於 torch.int8
。
使用這些計算,我們可以確定原始密集表示和新的稀疏表示的總記憶體佔用量。
這給了我們一個簡單的壓縮率公式,它僅取決於張量資料類型的位元寬度。
透過使用此公式,我們發現 torch.float16
或 torch.bfloat16
的壓縮率為 56.25%,而 torch.int8
的壓縮率為 62.5%。
建構稀疏半結構化張量¶
您可以使用 torch.to_sparse_semi_structured
函數將密集張量轉換為稀疏半結構化張量。
另請注意,我們僅支援 CUDA 張量,因為半結構化稀疏性的硬體相容性僅限於 NVIDIA GPU。
以下資料類型支援半結構化稀疏性。請注意,每種資料類型都有其自身的形狀約束和壓縮因子。
PyTorch dtype |
形狀約束 |
壓縮因子 |
稀疏模式 |
---|---|---|---|
|
張量必須是 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)
的輸出,而不是在張量以半結構化稀疏格式(例如這樣)具有 0 時使用 tensor
>>> a = torch.Tensor([0, 0, 1, 1]).tile((64, 16)).half().cuda()
>>> b = torch.rand(64, 64).half().cuda()
>>> c = torch.mm(a, b)
>>> a_sparse = to_sparse_semi_structured(a)
>>> torch.allclose(c, torch.mm(a_sparse, b))
True
使用半結構化稀疏性加速 nn.Linear¶
如果權重已經是半結構化稀疏的,您可以使用幾行程式碼來加速模型中的線性層
>>> input = torch.rand(64, 64).half().cuda()
>>> mask = torch.Tensor([0, 0, 1, 1]).tile((64, 16)).cuda().bool()
>>> linear = nn.Linear(64, 64).half().cuda()
>>> linear.weight = nn.Parameter(to_sparse_semi_structured(linear.weight.masked_fill(~mask, 0)))
稀疏 COO 張量¶
PyTorch 實作了所謂的座標格式(Coordinate format),或稱 COO 格式,作為實作稀疏張量(sparse tensor)的儲存格式之一。在 COO 格式中,指定的元素會儲存為元素索引和對應值的元組。特別是,
指定元素的索引會收集在大小為
(ndim, nse)
且元素類型為torch.int64
的indices
張量中,對應的值會收集在大小為
(nse,)
且具有任意整數或浮點數元素類型的values
張量中,
其中 ndim
是張量的維度,而 nse
是指定元素的數量。
注意
稀疏 COO 張量的記憶體消耗至少為 (ndim * 8 + <元素類型的大小,以位元組為單位>) * nse
位元組(加上儲存其他張量資料的常數 overhead)。
步幅張量(strided tensor)的記憶體消耗至少為 product(<張量形狀>) * <元素類型的大小,以位元組為單位>
。
例如,一個 10 000 x 10 000 的張量,具有 100 000 個非零的 32 位元浮點數,在使用 COO 張量佈局時,記憶體消耗至少為 (2 * 8 + 4) * 100 000 = 2 000 000
位元組,而使用預設的步幅張量佈局時,記憶體消耗為 10 000 * 10 000 * 4 = 400 000 000
位元組。 請注意,使用 COO 儲存格式可以節省 200 倍的記憶體。
建構¶
可以透過提供索引和值的兩個張量,以及稀疏張量的大小(當無法從索引和值張量推斷時)給函數 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) 維的張量,其在位置 (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
內容管理器實例來進行全域設定。預設情況下,稀疏張量不變量檢查是停用的。
未合併的稀疏 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 分別是批次 (batch)、稀疏 (sparse) 和稠密 (dense) 維度的數量,且滿足 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
是指定元素的數量。對於
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 維陣列。 此外,對於區塊稀疏壓縮佈局,2-D 區塊被視為正在指定的元素。 以一個 3 維區塊稀疏張量為例,它具有一個長度為 b
的批次維度,和一個形狀為 p, q
的區塊。 如果此張量有 n
個指定的元素,那麼實際上我們每個批次都指定了 n
個區塊。 此張量將具有形狀為 (b, n, p, q)
的 values
。 對於指定元素數量的這種解釋,來自於所有稀疏壓縮佈局都源自於對 2 維矩陣的壓縮。 批次維度被視為稀疏矩陣的堆疊,稠密維度將元素的含義從簡單的純量值更改為具有其自身維度的陣列。
稀疏 CSR 張量¶
CSR 格式相較於 COO 格式的主要優勢是更好地利用儲存空間,以及使用 MKL 和 MAGMA 後端進行稀疏矩陣-向量乘法等運算時速度更快。
在最簡單的情況下,(0 + 2 + 0) 維稀疏 CSR 張量由三個 1 維張量組成:crow_indices
、col_indices
和 values
。
crow_indices
張量由壓縮的列索引 (compressed row indices) 組成。 這是一個大小為nrows + 1
(列數加 1) 的 1 維張量。crow_indices
的最後一個元素是指定元素的數量nse
。 這個張量會根據給定列的起始位置,編碼values
和col_indices
中的索引。 張量中每個後續數字減去它之前的數字表示給定列中的元素數量。
col_indices
張量包含每個元素的行索引 (column indices)。 這是一個大小為nse
的 1 維張量。
values
張量包含 CSR 張量元素的值。 這是一個大小為nse
的 1 維張量。
注意
索引張量 crow_indices
和 col_indices
應該具有 torch.int64
(預設) 或 torch.int32
的元素類型。 如果您想使用啟用 MKL 的矩陣運算,請使用 torch.int32
。 這是由於 pytorch 的預設連結與 MKL LP64 相關,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 + <element type in bytes> * prod(densesize)) * nse) * prod(batchsize)
位元組(加上儲存其他張量資料的恆定開銷)。
使用與 稀疏 COO 格式介紹中的註解 相同的範例資料,一個具有 100,000 個非零 32 位浮點數的 10 000 x 10 000 張量的記憶體消耗至少為 (10000 * 8 + (8 + 4 * 1) * 100 000) * 1 = 1 280 000
位元組(當使用 CSR 張量佈局時)。請注意,與使用 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 張量建構 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
張量由壓縮的欄索引組成。這是一個 (B + 1)-D 張量,形狀為(*batchsize, ncols + 1)
。最後一個元素是指定的元素數量,nse
。這個張量根據給定欄位的起始位置,對values
和row_indices
中的索引進行編碼。張量中每個連續的數字減去它之前的數字表示給定欄位中元素的數量。
row_indices
張量包含每個元素的列索引。這是一個 (B + 1)-D 張量,形狀為(*batchsize, nse)
。
values
張量包含 CSC 張量元素的值。這是一個 (1 + K)-D 張量,形狀為(nse, *densesize)
。
CSC 張量的建構¶
稀疏 CSC 張量可以使用 torch.sparse_csc_tensor()
函數直接建構。使用者必須分別提供列索引、行索引和值張量,其中列索引必須使用 CSR 壓縮編碼指定。size
參數是可選的,如果不存在,將從 row_indices
和 ccol_indices
張量推導出來。
>>> ccol_indices = torch.tensor([0, 2, 4])
>>> row_indices = torch.tensor([0, 1, 0, 1])
>>> values = torch.tensor([1, 2, 3, 4])
>>> csc = torch.sparse_csc_tensor(ccol_indices, row_indices, values, dtype=torch.float64)
>>> csc
tensor(ccol_indices=tensor([0, 2, 4]),
row_indices=tensor([0, 1, 0, 1]),
values=tensor([1., 2., 3., 4.]), size=(2, 2), nnz=4,
dtype=torch.float64, layout=torch.sparse_csc)
>>> csc.to_dense()
tensor([[1., 3.],
[2., 4.]], dtype=torch.float64)
注意
稀疏 CSC 張量建構函式在列索引引數之前具有壓縮的欄索引引數。
可以使用 torch.Tensor.to_sparse_csc()
方法,從任何二維張量建構 (0 + 2 + 0) 維的稀疏 CSC 張量。 (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
張量由壓縮的列索引組成。 這是一個 (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 張量的建構¶
可以使用 torch.sparse_bsr_tensor()
函式直接建構稀疏 BSR 張量。 使用者必須分別提供列和欄區塊索引和值張量,其中列區塊索引必須使用 CSR 壓縮編碼來指定。size
引數是可選的,如果不存在,將從 crow_indices
和 col_indices
張量中推匯出來。
>>> crow_indices = torch.tensor([0, 2, 4])
>>> col_indices = torch.tensor([0, 1, 0, 1])
>>> values = torch.tensor([[[0, 1, 2], [6, 7, 8]],
... [[3, 4, 5], [9, 10, 11]],
... [[12, 13, 14], [18, 19, 20]],
... [[15, 16, 17], [21, 22, 23]]])
>>> bsr = torch.sparse_bsr_tensor(crow_indices, col_indices, values, dtype=torch.float64)
>>> bsr
tensor(crow_indices=tensor([0, 2, 4]),
col_indices=tensor([0, 1, 0, 1]),
values=tensor([[[ 0., 1., 2.],
[ 6., 7., 8.]],
[[ 3., 4., 5.],
[ 9., 10., 11.]],
[[12., 13., 14.],
[18., 19., 20.]],
[[15., 16., 17.],
[21., 22., 23.]]]),
size=(4, 6), nnz=4, dtype=torch.float64, layout=torch.sparse_bsr)
>>> bsr.to_dense()
tensor([[ 0., 1., 2., 3., 4., 5.],
[ 6., 7., 8., 9., 10., 11.],
[12., 13., 14., 15., 16., 17.],
[18., 19., 20., 21., 22., 23.]], dtype=torch.float64)
可以使用 torch.Tensor.to_sparse_bsr()
方法,從任何二維張量建構 (0 + 2 + 0) 維的稀疏 BSR 張量,該方法也需要指定值區塊大小
>>> dense = torch.tensor([[0, 1, 2, 3, 4, 5],
... [6, 7, 8, 9, 10, 11],
... [12, 13, 14, 15, 16, 17],
... [18, 19, 20, 21, 22, 23]])
>>> bsr = dense.to_sparse_bsr(blocksize=(2, 3))
>>> bsr
tensor(crow_indices=tensor([0, 2, 4]),
col_indices=tensor([0, 1, 0, 1]),
values=tensor([[[ 0, 1, 2],
[ 6, 7, 8]],
[[ 3, 4, 5],
[ 9, 10, 11]],
[[12, 13, 14],
[18, 19, 20]],
[[15, 16, 17],
[21, 22, 23]]]), size=(4, 6), nnz=4,
layout=torch.sparse_bsr)
稀疏 BSC 張量¶
稀疏 BSC (區塊壓縮稀疏欄) 張量格式實現了 BSC 格式,用於儲存二維張量,並擴展為支援批量稀疏 BSC 張量和作為多維張量區塊的值。
稀疏 BSC 張量由三個張量組成:ccol_indices
、row_indices
和 values
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 張量的建構¶
可以使用 torch.sparse_bsc_tensor()
函式直接建構稀疏 BSC 張量。 使用者必須分別提供列和欄區塊索引和值張量,其中欄區塊索引必須使用 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
參數。 以下範例說明了一種使用相同輸入資料,透過指定對應的 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)
支援的操作¶
線性代數操作¶
下表總結了稀疏矩陣上支援的線性代數操作,其中運算元的 layout 可能不同。此處 T[layout]
表示具有給定 layout 的張量。同樣地,M[layout]
表示矩陣(2-D PyTorch 張量),而 V[layout]
表示向量(1-D PyTorch 張量)。此外,f
表示純量(浮點數或 0-D PyTorch 張量),*
是逐元素乘法,而 @
是矩陣乘法。
PyTorch 操作 |
稀疏梯度? |
Layout 簽章 |
---|---|---|
否 |
|
|
否 |
|
|
否 |
|
|
否 |
|
|
否 |
|
|
否 |
|
|
否 |
|
|
否 |
|
|
否 |
|
|
是 |
|
|
否 |
|
|
否 |
|
|
否 |
|
|
否 |
|
|
否 |
|
|
否 |
|
|
是 |
|
|
否 |
|
|
否 |
|
|
否 |
|
|
是 |
|
|
是 |
|
其中「稀疏梯度?」欄位指示 PyTorch 操作是否支援對稀疏矩陣引數進行反向傳播。除了 torch.smm()
之外,所有 PyTorch 操作都支援對 strided 矩陣引數進行反向傳播。
注意
目前,PyTorch 不支援具有 layout 簽章 M[strided] @ M[sparse_coo]
的矩陣乘法。 但是,應用程式仍然可以使用矩陣關係 D @ S == (S.t() @ D.t()).t()
來計算它。
張量方法和稀疏¶
以下張量方法與稀疏張量有關
如果張量使用稀疏 COO 儲存 layout,則為 |
|
如果張量使用稀疏 CSR 儲存 layout,則為 |
|
返回稀疏張量 |
|
返回稀疏張量 |
|
返回一個新的稀疏張量,其值來自一個跨步張量 |
|
返回張量的稀疏副本。 |
|
將張量轉換為座標格式。 |
|
將張量轉換為壓縮列儲存格式(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 函式,專用於稀疏張量¶
使用給定的 |
|
使用給定的 |
|
使用給定的 |
|
使用給定的 |
|
使用給定的 |
|
使用給定的 |
|
回傳給定稀疏張量每一列的總和。 |
|
此函數在正向傳播 (forward) 的行為與 |
|
在 |
|
執行稀疏矩陣 |
|
將稀疏張量 |
|
執行 稀疏 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 對稀疏張量的支援。 |
零保留一元函數¶
我們的目標是支援所有「保留零的一元函數 (zero-preserving unary functions)」:將零映射到零的單一參數函數。
如果您發現我們缺少您需要的保留零的一元函數,請隨時提交 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()