torch.split

torch.split(tensor, split_size_or_sections, dim=0) [source]

Разбивает тензор на части. Каждая часть является представлением исходного тензора.

Если 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]]))