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 (Тензор) – тензор для разделения.
• split_size_or_sections (int) или (список(int)) – размер одной части или список размеров для каждой части
• dim (int) – измерение, по которому нужно разделить тензор.

Тип возвращаемого значения:

Список[Тензор]

Пример:

>>> 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]]))