torch.split

torch.split(tensor, split_size_or_sections, dim=0)

將張量分割成多個區塊。每個區塊都是原始張量的一個檢視 (view)。

如果 split_size_or_sections 是一個整數型別，則 tensor 將被分割成大小相等的區塊（如果可能的話）。 如果張量沿著給定的維度 dim 的大小不能被 split_size 整除，則最後一個區塊會比較小。

如果 split_size_or_sections 是一個列表，則 tensor 將被分割成 len(split_size_or_sections) 個區塊，其在 dim 維度上的大小將根據 split_size_or_sections 來決定。

參數

• tensor (Tensor) – 要分割的張量。

• split_size_or_sections (int) or (list(int)) – 單個區塊的大小，或是每個區塊大小的列表

• dim (int) – 要沿著分割張量的維度。

回傳型別

Tuple[Tensor, …]

範例

>>> a = torch.arange(10).reshape(5, 2)
>>> a
tensor([[0, 1],
        [2, 3],
        [4, 5],
        [6, 7],
        [8, 9]])
>>> torch.split(a, 2)
(tensor([[0, 1],
         [2, 3]]),
 tensor([[4, 5],
         [6, 7]]),
 tensor([[8, 9]]))
>>> torch.split(a, [1, 4])
(tensor([[0, 1]]),
 tensor([[2, 3],
         [4, 5],
         [6, 7],
         [8, 9]]))