pytorch

1. torch.flatten

将矩阵拉平

• 代码

import torch


# 随机创建一个矩阵，大小为(2,3,4)
a = torch.randint(0,10,size=(2,3,4))

# 将矩阵 a 按照第0维 ，拉平成一维向量，得到矩阵 b
b = torch.flatten(a)

# 将矩阵 a 按照第1维 ，拉平成二维向量，得到矩阵 c
c = torch.flatten(a,start_dim=1)

print(f"a={a}")
print(f"b={b}")

print(f"c={c}")

• 结果

a=tensor([[[0, 5, 1, 3],
         [4, 2, 8, 4],
         [7, 3, 0, 5]],

        [[2, 3, 8, 0],
         [1, 0, 8, 3],
         [8, 6, 6, 6]]])
b=tensor([0, 5, 1, 3, 4, 2, 8, 4, 7, 3, 0, 5, 2, 3, 8, 0, 1, 0, 8, 3, 8, 6, 6, 6])
c=tensor([[0, 5, 1, 3, 4, 2, 8, 4, 7, 3, 0, 5],
        [2, 3, 8, 0, 1, 0, 8, 3, 8, 6, 6, 6]])