Pytorch的API接口torch.flatten()

torch.flatten()

函数参数：

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

input: 展平的tensor

start_dim: 展平的起始维度

end_dim: 展平的结束维度

import torch

t = torch.tensor([[[1,2,3,4],
                   [5,6,7,8],
                   [1,5,9,7]],
                  [[3,5,7,8],
                   [2,5,8,9],
                   [4,5,6,7]],
                  [[6,0,5,8],
                   [9,6,4,7],
                   [5,8,1,6]]])
print(t.shape)

torch.Size([3, 3, 4])

可以理解成这个tensor的shape有3层，第一个值3指的是最外层中括号包含3个3 * 4的数组；

第二个3指的是每个3 * 4的数组有3个类似[1,2,3,4]；

第三个4指的是每个类似[1,2,3,4]中有4个元素；

参数start_dim

x = torch.flatten(t, 2)
print(x, x.shape)

tensor([[[1, 2, 3, 4],
         [5, 6, 7, 8],
         [1, 5, 9, 7]],

        [[3, 5, 7, 8],
         [2, 5, 8, 9],
         [4, 5, 6, 7]],

        [[6, 0, 5, 8],
         [9, 6, 4, 7],
         [5, 8, 1, 6]]]) torch.Size([3, 3, 4])

x = torch.flatten(t, 1)
print(x, x.shape)

tensor([[1, 2, 3, 4, 5, 6, 7, 8, 1, 5, 9, 7],
        [3, 5, 7, 8, 2, 5, 8, 9, 4, 5, 6, 7],
        [6, 0, 5, 8, 9, 6, 4, 7, 5, 8, 1, 6]]) torch.Size([3, 12])

x = torch.flatten(t, 0)
print(x, x.shape)

tensor([1, 2, 3, 4, 5, 6, 7, 8, 1, 5, 9, 7, 3, 5, 7, 8, 2, 5, 8, 9, 4, 5, 6, 7,
        6, 0, 5, 8, 9, 6, 4, 7, 5, 8, 1, 6]) torch.Size([36])

可以看出，start_dim是指定从哪一层开始展平，切记该参数是从0开始；

比如start_dim是2时，即从第3层展平，而第3层类似[1,2,3,4]已经是展平的；

start_dim是1时，即从第2层展平，即把[[1,2,3,4],[]…]展平成一维；

start_dim是0时，即从第1层展平，即把整个tensor展平成一维的

参数end_dim

y = torch.flatten(t, 0, 1)
print(y, y.shape)

tensor([[1, 2, 3, 4],
        [5, 6, 7, 8],
        [1, 5, 9, 7],
        [3, 5, 7, 8],
        [2, 5, 8, 9],
        [4, 5, 6, 7],
        [6, 0, 5, 8],
        [9, 6, 4, 7],
        [5, 8, 1, 6]]) torch.Size([9, 4])

y = torch.flatten(t, 0, 2)
print(y, y.shape)

tensor([1, 2, 3, 4, 5, 6, 7, 8, 1, 5, 9, 7, 3, 5, 7, 8, 2, 5, 8, 9, 4, 5, 6, 7,
        6, 0, 5, 8, 9, 6, 4, 7, 5, 8, 1, 6]) torch.Size([36])

理解了start_dim就容易理解end_dim了，start_dim是指定从哪一层开始展平，而end_dim是指定从哪一层开始结束展平；

end_dim默认值是-1，如果不指定，就展平到最后一层

refrence:

• https://pytorch.org/docs/master/generated/torch.flatten.html