Метод chunk
Метод chunk разделяет тензор на указанное количество частей
по заданному измерению. Метод принимает два основных параметра:
количество частей (chunks) и измерение (dim),
по которому производится разделение. Возвращается кортеж тензоров.
Синтаксис
t.chunk(chunks, dim=0)
Пример
Давайте разделим одномерный тензор на две части:
import torch
t = torch.tensor([1, 2, 3, 4, 5, 6])
res = t.chunk(2)
print(res)
Результат выполнения кода:
(tensor([1, 2, 3]), tensor([4, 5, 6]))
Пример
Разделим двумерный тензор по строкам (измерение 0):
import torch
t = torch.tensor([
[1, 2, 3],
[4, 5, 6],
[7, 8, 9],
[10, 11, 12],
])
res = t.chunk(2, dim=0)
print(res)
Результат выполнения кода:
(tensor([
[1, 2, 3],
[4, 5, 6],
]), tensor([
[7, 8, 9],
[10, 11, 12],
]))
Пример
Разделим тензор по столбцам (измерение 1):
import torch
t = torch.tensor([
[1, 2, 3, 4],
[5, 6, 7, 8],
])
res = t.chunk(2, dim=1)
print(res)
Результат выполнения кода:
(tensor([
[1, 2],
[5, 6],
]), tensor([
[3, 4],
[7, 8],
]))
Пример
Если размерность не делится на количество частей нацело, последняя часть будет меньше остальных:
import torch
t = torch.tensor([1, 2, 3, 4, 5])
res = t.chunk(3)
print(res)
Результат выполнения кода:
(tensor([1, 2]), tensor([3, 4]), tensor([5]))
