torch.flatten

torch.flatten(input, start_dim=0, end_dim=-1) → Tensor

将 input 通过重塑为一个一维张量来展平。如果传递了 start_dim 或 end_dim，则只展平从 start_dim 开始并以 end_dim 结束的维度。 input 中的元素顺序保持不变。

与 NumPy 的 flatten 不同（它总是复制 input 的数据），此函数可能会返回原始对象、视图或副本。如果没有展平任何维度，则返回原始对象 input。否则，如果 input 可以被视为展平后的形状，则返回该视图。最后，仅当 input 无法被视为展平后的形状时，才会复制 input 的数据。有关何时返回视图的详细信息，请参阅 torch.Tensor.view()。

注意

展平一个零维张量将返回一个一维视图。

参数

• input (Tensor) – 输入张量。

• start_dim (int) – 要展平的第一个维度

• end_dim (int) – 要展平的最后一个维度

示例

>>> t = torch.tensor([[[1, 2],
...                    [3, 4]],
...                   [[5, 6],
...                    [7, 8]]])
>>> torch.flatten(t)
tensor([1, 2, 3, 4, 5, 6, 7, 8])
>>> torch.flatten(t, start_dim=1)
tensor([[1, 2, 3, 4],
        [5, 6, 7, 8]])