torch.sparse¶
警告
稀疏張量的 PyTorch API 處於測試階段,在不久的將來可能會有所變更。我們非常歡迎您在 GitHub 問題中提出功能請求、錯誤報告和一般建議。
為什麼以及何時使用稀疏性¶
根據預設,PyTorch 將 torch.Tensor 元素連續儲存在實體記憶體中。這導致需要快速存取元素的各種陣列處理演算法能夠有效率地實作。
現在,有些使用者可能會決定使用 *元素大多為零值* 的張量來表示圖形鄰接矩陣、經過修剪的權重或點雲等資料。我們認識到這些都是重要的應用,並且目標是透過稀疏儲存格式為這些使用案例提供效能最佳化。
多年來,已經開發出各種稀疏儲存格式,例如 COO、CSR/CSC、半結構化、LIL 等。雖然它們的確切佈局不同,但它們都透過有效率地表示零值元素來壓縮資料。我們將未壓縮的值稱為 *已指定*,以區別於 *未指定* 的壓縮元素。
透過壓縮重複的零,稀疏儲存格式旨在節省各種 CPU 和 GPU 上的記憶體和計算資源。特別是對於高度稀疏或高度結構化的稀疏性,這可能會對效能產生顯著的影響。因此,稀疏儲存格式可以被視為一種效能最佳化。
與許多其他效能最佳化一樣,稀疏儲存格式並不總是有利的。當您嘗試將稀疏格式用於您的使用案例時,您可能會發現執行時間增加而不是減少。
如果您分析預期效能會大幅提升,但實際測量結果卻是下降,我們鼓勵您提出 GitHub 問題。這有助於我們優先實作有效的核心和更廣泛的效能最佳化。
我們讓您輕鬆嘗試不同的稀疏佈局,並在它們之間進行轉換,而無需考慮哪種佈局最適合您的特定應用。
功能概覽¶
我們希望透過為每個佈局提供轉換常式,可以直接從給定的密集張量建構稀疏張量。
在下一個範例中,我們將預設密集(跨步)佈局的二維張量轉換為由 COO 記憶體佈局支援的二維張量。在這種情況下,只會儲存非零元素的值和索引。
>>> 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 等裝置需要批次處理才能達到最佳效能,因此我們支援批次維度。
我們目前提供了一個非常簡單的批次處理版本,其中稀疏格式的每個組件本身都被批次處理。這也需要每個批次條目具有相同數量的指定元素。在此範例中,我們從三維(批次處理)密集張量建構一個三維(批次處理)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)
密集維度:另一方面,圖形嵌入等資料可能更適合視為向量的稀疏集合,而不是純量的集合。
在此範例中,我們從三維跨步張量建立一個具有 2 個稀疏維度和 1 個密集維度的三維混合 COO 張量。如果三維跨步張量中的一整列都是零,則不會儲存該列。但是,如果該列中的任何值是非零值,則會完整儲存這些值。這減少了索引的數量,因為我們每個列只需要一個索引,而不是每個元素一個索引。但它也增加了值的儲存量。因為只有 *完全* 為零的列才會被省略,而任何非零值元素的存在都會導致儲存整列。
>>> 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,而 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 ) 且元素類型為- torch.int16(如果 dtype 為 torch.float16 或 torch.bfloat16)或元素類型為- torch.int32(如果 dtype 為 torch.int8)的張量形式傳回 metadata_mask。
對於 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 函式,將稠密張量轉換為半結構化稀疏張量。
另請注意,由於半結構化稀疏性的硬體相容性僅限於 NVIDIA GPU,因此我們僅支援 CUDA 張量。
半結構化稀疏性支援以下資料類型。請注意,每種資料類型都有其形狀限制和壓縮因子。
| 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(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 實作了所謂的座標格式或 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 倍記憶體。
建構¶
您可以透過將索引和值的兩個張量,以及稀疏張量的大小(當無法從索引和值張量推斷出時)提供給函式 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 實例,為了將其與使用其他布局的*張量*實例區分開來,可以使用 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 成立。稀疏壓縮張量的稀疏維度數始終為兩個,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 維數組。同樣對於塊稀疏壓縮布局,二維塊被視為正在指定的元素。以壹個具有壹個長度為 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 預設連結到 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 + <元素類型大小(位元組)> * prod(densesize)) * nse) * prod(batchsize) 位元組(加上儲存其他張量資料的固定開銷)。
使用與 稀疏 COO 格式簡介中的註釋 相同的範例資料,當使用 CSR 張量佈局時,一個 10000 x 10000 的張量,其中包含 100000 個非零 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)-D 張量。
建構 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張量包含壓縮的列索引。這是一個形狀為(*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 張量¶
可以使用 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 編碼的所謂壓縮索引,以及與壓縮索引正交的所謂普通索引。這允許這些張量上的各種工具共享由張量佈局參數化的相同實現。
建構稀疏壓縮張量¶
可以使用 torch.sparse_compressed_tensor() 函數建構稀疏 CSR、CSC、BSR 和 CSC 張量,該函數具有與上述建構函數 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 表示標量(浮點數或零維 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 張量 的索引張量。 | 
| 當  | |
| 當  | 
以下 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 矩陣  | |
| 在  | |
| 執行稀疏矩陣  | |
| 將稀疏張量  | |
| 執行 稀疏 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()