pytorch中张量的拼接和分割

直接看代码，注意要举一反三

import torch

# 随机产生四个张量
a1 = torch.randn(3, 4)
a2 = torch.randn(3, 4)
a3 = torch.randn(3, 4)
a4 = torch.randn(3, 2)
print(a1)
print(a2)
print(a3)
print(a4)

print(torch.stack([a1, a2, a3], -1))  # 沿着最后一个维度即沿着列做堆叠，返回3*4*3的张量
print(torch.stack([a1,a2,a3],0)) # 沿着第一个维度即沿着行做堆叠，返回3*3*4的张量

print(torch.cat([a1,a2,a3,a4],-1)) # 沿着最后一个维度做拼接，返回3*14的张量


a = torch.randn(3, 6)  # 随机产生一个3*6的张量
print(a)

print(a.split([1, 2, 3], -1))  # 把张量沿着最后一个维度分割为三个张量，三个张量的维度分别为3*1,3*2,3*3
print(a.split([1, 2], 0))  # 把张量沿着第一个维度分割为两个张量，两个张量的维度分别为1*6,2*6

print(a.split(3, -1))  # 把张量沿着最后一个维度做分割，分割大小为3，输出的张量大小均为3*3
print(a.split(2, -1))  # 把张量沿着最后一个维度做分割，分割大小为2，输出的张量大小均为3*2

print(a.chunk(3,-1)) # 把张量沿着最后一个维度分割为3个张量，分割大小为2，输出的张量大小均为3*2
print(a.chunk(2,-1)) # 把张量沿着最后一个维度分割为2个张量，分割大小为2，输出的张量大小均为3*3

Done!!!