Python PyTorch flatten用法及代码示例

本文简要介绍python语言中 torch.flatten 的用法。

用法:

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

参数：

• input(Tensor) -输入张量。

• start_dim(int) -第一个要展平的dim

• end_dim(int) -最后一个dim要压平

通过将 input 重塑为一维张量来展平它。如果 start_dim 或 end_dim 已通过，则只有以 start_dim 开头并以 end_dim 结尾的尺寸会被展平。 input 中的元素顺序保持不变。

与 NumPy 的 flatten 不同，它总是复制输入的数据，这个函数可以返回原始对象、视图或副本。如果没有展平尺寸，则返回原始对象input。否则，如果输入可以被视为展平形状，则返回该视图。最后，只有当输入不能被视为扁平形状时，才会复制输入的数据。有关何时返回视图的详细信息，请参阅 torch.Tensor.view() 。

注意

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

例子：

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