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。否则，如果输入可以被视为展平后的形状，则返回该视图。只有在无法将输入视为展平时，才会复制输入的数据。有关何时返回视图的详细信息，请参见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]])