torch.squeeze

torch.squeeze(input, dim=None) → Tensor

Возвращает тензор, из которого удалены все измерения размера input.

Например, если input имеет форму: $(A \times 1 \times B \times C \times 1 \times D)$, то возвращаемый out тензор будет иметь форму: $(A \times B \times C \times D)$.

Если задан параметр dim, операция сжатия выполняется только по указанному измерению. Если input имеет форму $(A \times 1 \times B)$, squeeze(input, 0) оставляет тензор неизменным, но squeeze(input, 1) сжимает тензор до формы $(A \times B)$.

Примечание

Возвращаемый тензор использует ту же память, что и входной тензор, поэтому изменение содержимого одного из них повлияет на содержимое другого.

Предупреждение

Если тензор имеет размерность пакета (batch) размером 1, то squeeze(input) также удалит размерность пакета, что может привести к неожиданным ошибкам.

Параметры:

• input (Tensor) – входной тензор.
• dim (int, необязательно) – если задано, входной тензор будет сжат только по этому измерению

Пример:

>>> x = torch.zeros(2, 1, 2, 1, 2)
>>> x.size()
torch.Size([2, 1, 2, 1, 2])
>>> y = torch.squeeze(x)
>>> y.size()
torch.Size([2, 2, 2])
>>> y = torch.squeeze(x, 0)
>>> y.size()
torch.Size([2, 1, 2, 1, 2])
>>> y = torch.squeeze(x, 1)
>>> y.size()
torch.Size([2, 2, 1, 2])