快捷鍵

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 目前支援 COOCSRCSCBSRBSC

我們還有一個原型實作來支援 :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.shapee = bitwidth(tensor.dtype),因此 e = 16 對於 torch.float16torch.bfloat16,以及 e = 8 對於 torch.int8

Mdense=r×c×eMsparse=Mspecified+Mmetadata=r×c2×e+r×c2×2=rce2+rc=rce(12+1e)M_{dense} = r \times c \times e \\ M_{sparse} = M_{specified} + M_{metadata} = r \times \frac{c}{2} \times e + r \times \frac{c}{2} \times 2 = \frac{rce}{2} + rc =rce(\frac{1}{2} +\frac{1}{e})

使用這些計算,我們可以確定原始密集表示和新的稀疏表示的總記憶體佔用量。

這給了我們一個簡單的壓縮率公式,它僅取決於張量資料類型的位元寬度。

C=MsparseMdense=12+1eC = \frac{M_{sparse}}{M_{dense}} = \frac{1}{2} + \frac{1}{e}

透過使用此公式,我們發現 torch.float16torch.bfloat16 的壓縮率為 56.25%,而 torch.int8 的壓縮率為 62.5%。

建構稀疏半結構化張量

您可以使用 torch.to_sparse_semi_structured 函數將密集張量轉換為稀疏半結構化張量。

另請注意,我們僅支援 CUDA 張量,因為半結構化稀疏性的硬體相容性僅限於 NVIDIA GPU。

以下資料類型支援半結構化稀疏性。請注意,每種資料類型都有其自身的形狀約束和壓縮因子。

PyTorch dtype

形狀約束

壓縮因子

稀疏模式

torch.float16

張量必須是 2D,且 (r, c) 必須都是 64 的正倍數

9/16

2:4

torch.bfloat16

張量必須是 2D,且 (r, c) 必須都是 64 的正倍數

9/16

2:4

torch.int8

張量必須是 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.int64indices 張量中,

  • 對應的值會收集在大小為 (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.int64indices 張量中,

  • 對應的(張量)值會收集在大小為 (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 指定多個值,34,這會產生一個一維的未合併張量

>>> 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() 方法的輸出是一個具有以下屬性的稀疏張量

注意

在大多數情況下,您不必關心稀疏張量是否合併,因為大多數操作對於合併或未合併的稀疏張量的工作方式都相同。

但是,某些操作在未合併的張量上可以更有效地實現,而某些操作在合併的張量上可以更有效地實現。

例如,稀疏 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_sparsetorch.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_size0 <= 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_indicescol_indicesvalues

  • crow_indices 張量由壓縮的列索引 (compressed row indices) 組成。 這是一個大小為 nrows + 1 (列數加 1) 的 1 維張量。 crow_indices 的最後一個元素是指定元素的數量 nse。 這個張量會根據給定列的起始位置,編碼 valuescol_indices 中的索引。 張量中每個後續數字減去它之前的數字表示給定列中的元素數量。

  • col_indices 張量包含每個元素的行索引 (column indices)。 這是一個大小為 nse 的 1 維張量。

  • values 張量包含 CSR 張量元素的值。 這是一個大小為 nse 的 1 維張量。

注意

索引張量 crow_indicescol_indices 應該具有 torch.int64 (預設) 或 torch.int32 的元素類型。 如果您想使用啟用 MKL 的矩陣運算,請使用 torch.int32。 這是由於 pytorch 的預設連結與 MKL LP64 相關,MKL LP64 使用 32 位元整數索引。

在一般情況下,(B + 2 + K) 維稀疏 CSR 張量由兩個 (B + 1) 維索引張量 crow_indicescol_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) == Blen(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_indicescol_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_indicesrow_indicesvalues

  • ccol_indices 張量由壓縮的欄索引組成。這是一個 (B + 1)-D 張量,形狀為 (*batchsize, ncols + 1)。最後一個元素是指定的元素數量,nse。這個張量根據給定欄位的起始位置,對 valuesrow_indices 中的索引進行編碼。張量中每個連續的數字減去它之前的數字表示給定欄位中元素的數量。

  • row_indices 張量包含每個元素的列索引。這是一個 (B + 1)-D 張量,形狀為 (*batchsize, nse)

  • values 張量包含 CSC 張量元素的值。這是一個 (1 + K)-D 張量,形狀為 (nse, *densesize)

CSC 張量的建構

稀疏 CSC 張量可以使用 torch.sparse_csc_tensor() 函數直接建構。使用者必須分別提供列索引、行索引和值張量,其中列索引必須使用 CSR 壓縮編碼指定。size 參數是可選的,如果不存在,將從 row_indicesccol_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_indicescol_indicesvalues

  • crow_indices 張量由壓縮的列索引組成。 這是一個 (B + 1) 維張量,形狀為 (*batchsize, nrowblocks + 1)。 最後一個元素是指定的區塊數,nse。 這個張量根據給定的列區塊的起始位置對 valuescol_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_indicescol_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_indicesrow_indicesvalues

  • ccol_indices 張量由壓縮的欄索引組成。 這是一個 (B + 1) 維張量,形狀為 (*batchsize, ncolblocks + 1)。 最後一個元素是指定的區塊數,nse。 這個張量根據給定的列區塊的起始位置對 valuesrow_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_indicesrow_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 簽章

torch.mv()

M[sparse_coo] @ V[strided] -> V[strided]

torch.mv()

M[sparse_csr] @ V[strided] -> V[strided]

torch.matmul()

M[sparse_coo] @ M[strided] -> M[strided]

torch.matmul()

M[sparse_csr] @ M[strided] -> M[strided]

torch.matmul()

M[SparseSemiStructured] @ M[strided] -> M[strided]

torch.matmul()

M[strided] @ M[SparseSemiStructured] -> M[strided]

torch.mm()

M[sparse_coo] @ M[strided] -> M[strided]

torch.mm()

M[SparseSemiStructured] @ M[strided] -> M[strided]

torch.mm()

M[strided] @ M[SparseSemiStructured] -> M[strided]

torch.sparse.mm()

M[sparse_coo] @ M[strided] -> M[strided]

torch.smm()

M[sparse_coo] @ M[strided] -> M[sparse_coo]

torch.hspmm()

M[sparse_coo] @ M[strided] -> M[hybrid sparse_coo]

torch.bmm()

T[sparse_coo] @ T[strided] -> T[strided]

torch.addmm()

f * M[strided] + f * (M[sparse_coo] @ M[strided]) -> M[strided]

torch.addmm()

f * M[strided] + f * (M[SparseSemiStructured] @ M[strided]) -> M[strided]

torch.addmm()

f * M[strided] + f * (M[strided] @ M[SparseSemiStructured]) -> M[strided]

torch.sparse.addmm()

f * M[strided] + f * (M[sparse_coo] @ M[strided]) -> M[strided]

torch.sparse.spsolve()

SOLVE(M[sparse_csr], V[strided]) -> V[strided]

torch.sspaddmm()

f * M[sparse_coo] + f * (M[sparse_coo] @ M[strided]) -> M[sparse_coo]

torch.lobpcg()

GENEIG(M[sparse_coo]) -> M[strided], M[strided]

torch.pca_lowrank()

PCA(M[sparse_coo]) -> M[strided], M[strided], M[strided]

torch.svd_lowrank()

SVD(M[sparse_coo]) -> M[strided], M[strided], M[strided]

其中「稀疏梯度?」欄位指示 PyTorch 操作是否支援對稀疏矩陣引數進行反向傳播。除了 torch.smm() 之外,所有 PyTorch 操作都支援對 strided 矩陣引數進行反向傳播。

注意

目前,PyTorch 不支援具有 layout 簽章 M[strided] @ M[sparse_coo] 的矩陣乘法。 但是,應用程式仍然可以使用矩陣關係 D @ S == (S.t() @ D.t()).t() 來計算它。

張量方法和稀疏

以下張量方法與稀疏張量有關

Tensor.is_sparse

如果張量使用稀疏 COO 儲存 layout,則為 True,否則為 False

Tensor.is_sparse_csr

如果張量使用稀疏 CSR 儲存 layout,則為 True,否則為 False

Tensor.dense_dim

返回稀疏張量 self 中的稠密維度數量。

Tensor.sparse_dim

返回稀疏張量 self 中的稀疏維度數量。

Tensor.sparse_mask

返回一個新的稀疏張量,其值來自一個跨步張量 self,該張量根據稀疏張量 mask 的索引進行過濾。

Tensor.to_sparse

返回張量的稀疏副本。

Tensor.to_sparse_coo

將張量轉換為座標格式

Tensor.to_sparse_csr

將張量轉換為壓縮列儲存格式(CSR)。

Tensor.to_sparse_csc

將張量轉換為壓縮行儲存格式(CSC)。

Tensor.to_sparse_bsr

將張量轉換為給定區塊大小的區塊稀疏列(BSR)儲存格式。

Tensor.to_sparse_bsc

將張量轉換為給定區塊大小的區塊稀疏行(BSC)儲存格式。

Tensor.to_dense

如果 self 不是跨步張量,則建立 self 的跨步副本,否則返回 self

Tensor.values

返回 稀疏 COO 張量 的值張量。

以下 Tensor 方法專用於稀疏 COO 張量

Tensor.coalesce

如果 self 是一個未合併的張量,則返回 self 的合併副本。

Tensor.sparse_resize_

self 稀疏張量 調整為所需的大小以及稀疏和密集維度的數量。

Tensor.sparse_resize_and_clear_

稀疏張量 self 中刪除所有指定的元素,並將 self 調整為所需的大小以及稀疏和密集維度的數量。

Tensor.is_coalesced

如果 self 是合併的 稀疏 COO 張量,則返回 True,否則返回 False

Tensor.indices

返回 稀疏 COO 張量 的索引張量。

以下方法專用於稀疏 CSR 張量稀疏 BSR 張量

Tensor.crow_indices

self 是佈局為 sparse_csr 的稀疏 CSR 張量時,返回包含 self 張量的壓縮列索引的張量。

Tensor.col_indices

self 是佈局為 sparse_csr 的稀疏 CSR 張量時,返回包含 self 張量的行索引的張量。

以下方法專用於稀疏 CSC 張量稀疏 BSC 張量

Tensor.row_indices

Tensor.ccol_indices

以下 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 函式,專用於稀疏張量

sparse_coo_tensor

使用給定的 indices,建構一個在 COO(rdinate) 格式的 稀疏張量,並指定其值。

sparse_csr_tensor

使用給定的 crow_indicescol_indices,建構一個在 CSR (Compressed Sparse Row, 壓縮稀疏列) 格式的 稀疏張量,並指定其值。

sparse_csc_tensor

使用給定的 ccol_indicesrow_indices,建構一個在 CSC (Compressed Sparse Column, 壓縮稀疏行) 格式的 稀疏張量,並指定其值。

sparse_bsr_tensor

使用給定的 crow_indicescol_indices,建構一個在 BSR (Block Compressed Sparse Row, 區塊壓縮稀疏列) 格式的 稀疏張量,並指定給定的二維區塊。

sparse_bsc_tensor

使用給定的 ccol_indicesrow_indices,建構一個在 BSC (Block Compressed Sparse Column, 區塊壓縮稀疏行) 格式的 稀疏張量,並指定給定的二維區塊。

sparse_compressed_tensor

使用給定的 compressed_indicesplain_indices,建構一個在壓縮稀疏格式 (Compressed Sparse format) - CSR、CSC、BSR 或 BSC - 的 稀疏張量,並指定其值。

sparse.sum

回傳給定稀疏張量每一列的總和。

sparse.addmm

此函數在正向傳播 (forward) 的行為與 torch.addmm() 完全相同,但它支援稀疏 COO 矩陣 mat1 的反向傳播 (backward)。

sparse.sampled_addmm

input 的稀疏模式 (sparsity pattern) 所指定的位置,執行稠密矩陣 mat1mat2 的矩陣乘法。

sparse.mm

執行稀疏矩陣 mat1 的矩陣乘法。

sspaddmm

將稀疏張量 mat1 與稠密張量 mat2 相乘,然後將稀疏張量 input 加到結果中。

hspmm

執行 稀疏 COO 矩陣 mat1 與跨步 (strided) 矩陣 mat2 的矩陣乘法。

smm

執行稀疏矩陣 input 與稠密矩陣 mat 的矩陣乘法。

sparse.softmax

應用 softmax 函數。

sparse.spsolve

計算具有唯一解的線性方程平方系統的解。

sparse.log_softmax

應用 softmax 函數,然後取對數。

sparse.spdiags

藉由將 diagonals 的列中的值沿著輸出的指定對角線放置,來建立稀疏 2D 張量。

其他函數

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

若要管理檢查稀疏張量不變性,請參閱

sparse.check_sparse_tensor_invariants

一個用來控制檢查稀疏張量不變性的工具。

若要將稀疏張量與 gradcheck() 函數一起使用,請參閱

sparse.as_sparse_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()

文件

存取 PyTorch 的完整開發人員文件

檢視文件

教學課程

取得適用於初學者和進階開發人員的深入教學課程

檢視教學課程

資源

尋找開發資源並獲得您的問題解答

檢視資源