捷徑

torch.nested

簡介

警告

巢狀張量的 PyTorch API 處於原型階段,並將在不久的將來發生變化。

巢狀張量允許使用者將張量列表打包成單個、高效的資料結構。

對輸入張量的唯一限制是它們的維度必須匹配。

這使得更有效的元資料表示和對專用內核的訪問成為可能。

巢狀張量的一種應用是在各個領域表達順序資料。雖然傳統方法是填充可變長度序列,但巢狀張量使用戶能夠繞過填充。在巢狀張量上呼叫操作的 API 與常規 torch.Tensor 的 API 沒有什麼不同,這應該允許與現有模型無縫集成,主要區別在於 輸入的建構

由於這是一個原型功能,因此 支援的操作 仍然有限。但是,我們歡迎問題、功能請求和貢獻。有關貢獻的更多資訊可以在 此自述文件中 找到。

建構

建構很簡單,只需將張量列表傳遞給 torch.nested.nested_tensor 建構函數即可。

>>> a, b = torch.arange(3), torch.arange(5) + 3
>>> a
tensor([0, 1, 2])
>>> b
tensor([3, 4, 5, 6, 7])
>>> nt = torch.nested.nested_tensor([a, b])
>>> nt
nested_tensor([
  tensor([0, 1, 2]),
    tensor([3, 4, 5, 6, 7])
    ])

資料類型、裝置以及是否需要梯度可以透過通常的關鍵字參數來選擇。

>>> nt = torch.nested.nested_tensor([a, b], dtype=torch.float32, device="cuda", requires_grad=True)
>>> nt
nested_tensor([
  tensor([0., 1., 2.], device='cuda:0', requires_grad=True),
  tensor([3., 4., 5., 6., 7.], device='cuda:0', requires_grad=True)
], device='cuda:0', requires_grad=True)

torch.as_tensor 類似,torch.nested.as_nested_tensor 可用於保留從傳遞給建構函數的張量的自動梯度歷史記錄。有關更多資訊,請參閱關於 巢狀張量建構函數和轉換函數 的章節。

為了形成有效的巢狀張量,所有傳遞的張量都需要在維度上匹配,但其他屬性則不需要。

>>> a = torch.randn(3, 50, 70) # image 1
>>> b = torch.randn(3, 128, 64) # image 2
>>> nt = torch.nested.nested_tensor([a, b], dtype=torch.float32)
>>> nt.dim()
4

如果其中一個維度不匹配,建構函數會引發錯誤。

>>> a = torch.randn(50, 128) # text 1
>>> b = torch.randn(3, 128, 64) # image 2
>>> nt = torch.nested.nested_tensor([a, b], dtype=torch.float32)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
RuntimeError: All Tensors given to nested_tensor must have the same dimension. Found dimension 3 for Tensor at index 1 and dimension 2 for Tensor at index 0.

請注意,傳遞的張量將被複製到一個連續的記憶體片段中。生成的巢狀張量會分配新的記憶體來儲存它們,並且不會保留引用。

目前我們只支援一級巢狀,即一個簡單的、扁平的張量列表。將來我們可以添加對多級巢狀的支援,例如完全由張量列表組成的列表。請注意,對於此擴展,重要的是在條目之間保持均勻的巢狀級別,以便生成的巢狀張量具有明確的維度。如果您需要此功能,請鼓勵您提出功能請求,以便我們可以追蹤它並相應地進行規劃。

大小

儘管巢狀張量不支援 .size()(或 .shape),但如果維度 i 是規則的,則它支援 .size(i)

>>> a = torch.randn(50, 128) # text 1
>>> b = torch.randn(32, 128) # text 2
>>> nt = torch.nested.nested_tensor([a, b], dtype=torch.float32)
>>> nt.size(0)
2
>>> nt.size(1)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
RuntimeError: Given dimension 1 is irregular and does not have a size.
>>> nt.size(2)
128

如果所有維度都是規則的,則巢狀張量在語義上應該與常規的 torch.Tensor 沒有區別。

>>> a = torch.randn(20, 128) # text 1
>>> nt = torch.nested.nested_tensor([a, a], dtype=torch.float32)
>>> nt.size(0)
2
>>> nt.size(1)
20
>>> nt.size(2)
128
>>> torch.stack(nt.unbind()).size()
torch.Size([2, 20, 128])
>>> torch.stack([a, a]).size()
torch.Size([2, 20, 128])
>>> torch.equal(torch.stack(nt.unbind()), torch.stack([a, a]))
True

將來,我們可能會使其更容易檢測到這種情況並無縫轉換。

如果您需要此功能(或任何其他相關功能),請提出功能請求。

解綁

unbind 允許您檢索成分的視圖。

>>> import torch
>>> a = torch.randn(2, 3)
>>> b = torch.randn(3, 4)
>>> nt = torch.nested.nested_tensor([a, b], dtype=torch.float32)
>>> nt
nested_tensor([
  tensor([[ 1.2286, -1.2343, -1.4842],
          [-0.7827,  0.6745,  0.0658]]),
  tensor([[-1.1247, -0.4078, -1.0633,  0.8083],
          [-0.2871, -0.2980,  0.5559,  1.9885],
          [ 0.4074,  2.4855,  0.0733,  0.8285]])
])
>>> nt.unbind()
(tensor([[ 1.2286, -1.2343, -1.4842],
        [-0.7827,  0.6745,  0.0658]]), tensor([[-1.1247, -0.4078, -1.0633,  0.8083],
        [-0.2871, -0.2980,  0.5559,  1.9885],
        [ 0.4074,  2.4855,  0.0733,  0.8285]]))
>>> nt.unbind()[0] is not a
True
>>> nt.unbind()[0].mul_(3)
tensor([[ 3.6858, -3.7030, -4.4525],
        [-2.3481,  2.0236,  0.1975]])
>>> nt
nested_tensor([
  tensor([[ 3.6858, -3.7030, -4.4525],
          [-2.3481,  2.0236,  0.1975]]),
  tensor([[-1.1247, -0.4078, -1.0633,  0.8083],
          [-0.2871, -0.2980,  0.5559,  1.9885],
          [ 0.4074,  2.4855,  0.0733,  0.8285]])
])

請注意,nt.unbind()[0] 不是副本,而是底層記憶體的切片,它表示巢狀張量的第一個條目或成分。

巢狀張量建構函數和轉換函數

以下函數與巢狀張量相關

torch.nested.nested_tensor(tensor_list, *, dtype=None, layout=None, device=None, requires_grad=False, pin_memory=False)[原始碼]

從張量列表 tensor_list 建構一個沒有自動梯度歷史記錄的巢狀張量(也稱為“葉張量”,請參閱 Autograd 機制)。

參數
  • tensor_list (List[array_like]) – 張量列表,或任何可以傳遞給 torch.tensor 的內容,

  • 維度。 (其中列表的每個元素都具有相同的) –

關鍵字參數
  • dtype (torch.dtype,可選) – 返回巢狀張量的所需類型。預設值:如果為 None,則與清單中最左側張量的 torch.dtype 相同。

  • layout (torch.layout,可選) – 返回巢狀張量的所需佈局。僅支援跨步和交錯佈局。預設值:如果為 None,則為跨步佈局。

  • device (torch.device,可選) – 返回巢狀張量的所需設備。預設值:如果為 None,則與清單中最左側張量的 torch.device 相同

  • requires_grad (bool可選) – 自動求導是否應記錄返回巢狀張量上的操作。預設值:False

  • pin_memory (bool可選) – 如果設定,則返回的巢狀張量將分配在固定內存中。僅適用於 CPU 張量。預設值:False

回傳型別

張量

範例

>>> a = torch.arange(3, dtype=torch.float, requires_grad=True)
>>> b = torch.arange(5, dtype=torch.float, requires_grad=True)
>>> nt = torch.nested.nested_tensor([a, b], requires_grad=True)
>>> nt.is_leaf
True
torch.nested.as_nested_tensor(ts, dtype=None, device=None, layout=None)[原始碼]

從張量或張量的清單/元組建構一個保留自動求導歷史記錄的巢狀張量。

如果傳遞巢狀張量,則將直接返回它,除非設備/資料類型/佈局不同。請注意,轉換設備/資料類型將導致複製,而此函式目前不支援轉換佈局。

如果傳遞非巢狀張量,則將其視為一批大小一致的組成部分。如果傳遞的設備/資料類型與輸入的不同,或者如果輸入是非連續的,則將產生副本。否則,將直接使用輸入的儲存空間。

如果提供張量清單,則在建構巢狀張量期間,清單中的張量將始終被複製。

參數

ts (張量清單[張量] 或 元組[張量]) – 要視為巢狀張量的張量,或具有相同 ndim 的張量清單/元組

關鍵字參數
  • dtype (torch.dtype,可選) – 返回巢狀張量的所需類型。預設值:如果為 None,則與清單中最左側張量的 torch.dtype 相同。

  • device (torch.device,可選) – 返回巢狀張量的所需設備。預設值:如果為 None,則與清單中最左側張量的 torch.device 相同

  • layout (torch.layout,可選) – 返回巢狀張量的所需佈局。僅支援跨步和交錯佈局。預設值:如果為 None,則為跨步佈局。

回傳型別

張量

範例

>>> a = torch.arange(3, dtype=torch.float, requires_grad=True)
>>> b = torch.arange(5, dtype=torch.float, requires_grad=True)
>>> nt = torch.nested.as_nested_tensor([a, b])
>>> nt.is_leaf
False
>>> fake_grad = torch.nested.nested_tensor([torch.ones_like(a), torch.zeros_like(b)])
>>> nt.backward(fake_grad)
>>> a.grad
tensor([1., 1., 1.])
>>> b.grad
tensor([0., 0., 0., 0., 0.])
>>> c = torch.randn(3, 5, requires_grad=True)
>>> nt2 = torch.nested.as_nested_tensor(c)
torch.nested.to_padded_tensor(input, padding, output_size=None, out=None) 張量

通過填充 input 巢狀張量返回一個新的(非巢狀)張量。前導條目將填充巢狀資料,而尾隨條目將填充。

警告

to_padded_tensor() 總是複製底層資料,因為巢狀和非巢狀張量在內存佈局上有所不同。

參數

padding (float) – 尾隨條目的填充值。

關鍵字參數
  • output_size (元組[int]) – 輸出張量的大小。如果給定,則它必須足夠大以包含所有巢狀資料;否則,將通過沿每個維度採用每個巢狀子張量的最大大小來推斷。

  • out (張量可選) – 輸出張量。

範例

>>> nt = torch.nested.nested_tensor([torch.randn((2, 5)), torch.randn((3, 4))])
nested_tensor([
  tensor([[ 1.6862, -1.1282,  1.1031,  0.0464, -1.3276],
          [-1.9967, -1.0054,  1.8972,  0.9174, -1.4995]]),
  tensor([[-1.8546, -0.7194, -0.2918, -0.1846],
          [ 0.2773,  0.8793, -0.5183, -0.6447],
          [ 1.8009,  1.8468, -0.9832, -1.5272]])
])
>>> pt_infer = torch.nested.to_padded_tensor(nt, 0.0)
tensor([[[ 1.6862, -1.1282,  1.1031,  0.0464, -1.3276],
         [-1.9967, -1.0054,  1.8972,  0.9174, -1.4995],
         [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000]],
        [[-1.8546, -0.7194, -0.2918, -0.1846,  0.0000],
         [ 0.2773,  0.8793, -0.5183, -0.6447,  0.0000],
         [ 1.8009,  1.8468, -0.9832, -1.5272,  0.0000]]])
>>> pt_large = torch.nested.to_padded_tensor(nt, 1.0, (2, 4, 6))
tensor([[[ 1.6862, -1.1282,  1.1031,  0.0464, -1.3276,  1.0000],
         [-1.9967, -1.0054,  1.8972,  0.9174, -1.4995,  1.0000],
         [ 1.0000,  1.0000,  1.0000,  1.0000,  1.0000,  1.0000],
         [ 1.0000,  1.0000,  1.0000,  1.0000,  1.0000,  1.0000]],
        [[-1.8546, -0.7194, -0.2918, -0.1846,  1.0000,  1.0000],
         [ 0.2773,  0.8793, -0.5183, -0.6447,  1.0000,  1.0000],
         [ 1.8009,  1.8468, -0.9832, -1.5272,  1.0000,  1.0000],
         [ 1.0000,  1.0000,  1.0000,  1.0000,  1.0000,  1.0000]]])
>>> pt_small = torch.nested.to_padded_tensor(nt, 2.0, (2, 2, 2))
RuntimeError: Value in output_size is less than NestedTensor padded size. Truncation is not supported.

支援的運算

在本節中,我們總結了 NestedTensor 當前支援的運算及其具有的任何限制。

PyTorch 運算

限制

torch.matmul()

支援兩個(>= 3d)巢狀張量之間的矩陣乘法,其中最後兩個維度是矩陣維度,並且前導(批次)維度的大小相同(即,尚不支援批次維度的廣播)。

torch.bmm()

支援兩個 3 維巢狀張量的批次矩陣乘法。

torch.nn.Linear()

支援 3 維巢狀輸入和密集的 2 維權重矩陣。

torch.nn.functional.softmax()

支援除 dim=0 之外的所有維度上的 softmax。

torch.nn.Dropout()

行為與在常規張量上的行為相同。

torch.Tensor.masked_fill()

行為與在常規張量上的行為相同。

torch.relu()

行為與在常規張量上的行為相同。

torch.gelu()

行為與在常規張量上的行為相同。

torch.silu()

行為與在常規張量上的行為相同。

torch.abs()

行為與在常規張量上的行為相同。

torch.sgn()

行為與在常規張量上的行為相同。

torch.logical_not()

行為與在常規張量上的行為相同。

torch.neg()

行為與在常規張量上的行為相同。

torch.sub()

支援兩個巢狀張量的逐元素減法。

torch.add()

支援兩個巢狀張量的逐元素加法。支援將標量添加到巢狀張量。

torch.mul()

支援兩個巢狀張量的逐元素乘法。支援將巢狀張量乘以標量。

torch.select()

支援沿所有維度進行選擇。

torch.clone()

行為與在常規張量上的行為相同。

torch.detach()

行為與在常規張量上的行為相同。

torch.unbind()

僅支援沿 dim=0 取消綁定。

torch.reshape()

支援在保留 dim=0 的大小的情況下進行整形(即,不能更改巢狀張量的數量)。與常規張量不同,這裡的大小 -1 表示繼承現有大小。特別是,不規則維度的唯一有效大小是 -1。大小推斷尚未實作,因此對於新維度,大小不能為 -1

torch.Tensor.reshape_as()

reshape 的限制類似。

torch.transpose()

支援除 dim=0 之外的所有維度的轉置。

torch.Tensor.view()

新形狀的規則與 reshape 的規則類似。

torch.empty_like()

行為類似於常規張量;返回一個與輸入的巢狀結構匹配的新空巢狀張量(即具有未初始化的值)。

torch.randn_like()

行為類似於常規張量;返回一個新的巢狀張量,其值根據與輸入的巢狀結構匹配的標準正態分佈隨機初始化。

torch.zeros_like()

行為類似於常規張量;返回一個與輸入的巢狀結構匹配的所有值為零的新巢狀張量。

torch.nn.LayerNorm()

normalized_shape 參數被限制為不擴展到 NestedTensor 的不規則維度。

文件

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

檢視文件

教學課程

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

檢視教學課程

資源

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

檢視資源