pytorch: tensor的拆分和拼接（split, chunk, cat, stack）

拆分

torch.split[按数量拆分]

torch.split() 按数量进行拆分：

import torch

x = torch.rand(8, 16, 3, 7,7)

y1 = torch.split(x, 8, dim=1)  # 第二维度拆分成大小为8
for y in y1:
    print(y.shape)

out:
torch.Size([8, 8, 3, 7, 7])
torch.Size([8, 8, 3, 7, 7])

当不能均等份拆分时，剩余的数量作为张量（tensor）:

y2 = torch.split(x, 5, dim=1)  # 第二维度拆分为5, 剩余不足5的数量作为tensor
for y in y2:
    print(y.shape)

out:
torch.Size([8, 5, 3, 7, 7])
torch.Size([8, 5, 3, 7, 7])
torch.Size([8, 5, 3, 7, 7])
torch.Size([8, 1, 3, 7, 7])

torch.chunk[按块拆分]

torch.chunk() 按块进行拆分：

y3 = torch.chunk(x, 2, dim=1)  # 第二维度拆分成2块
for y in y3:
    print(y.shape)

out:
torch.Size([8, 8, 3, 7, 7])
torch.Size([8, 8, 3, 7, 7])

当不能均等份拆分时，拆分的维度大小为=N//n + 1，其中，N为张量x的待拆分的维度16，n为拆分成块的数量，则拆分成每块的数量为：16//3 + 1= 5+1=6，剩余不足6的作为张量，如：

y4 = torch.chunk(x, 3, dim=1)  # 第二维度拆分为3块
for y in y4:
    print(y.shape)

out:
torch.Size([8, 6, 3, 7, 7])
torch.Size([8, 6, 3, 7, 7])
torch.Size([8, 4, 3, 7, 7])

拼接

torch.cat[按原有维度拼接]

torch.cat() 在原有的维度进行拼接，其中，进行拼接的维度大小可以不一样，其余维度大小必须一样，如：

z1 = torch.cat(y1, dim=1)  # 基于第二维度合并
print(z1.shape)
z2 = torch.cat(y2, dim=1)
print(z2.shape)

out:
torch.Size([8, 16, 3, 7, 7])
torch.Size([8, 16, 3, 7, 7])
y1所有成员数量为2，且第二维度均为8，拼接后第二维度为16；
y2的成员为数量为4，虽然第二维度不一致，但基于第二维度拼接，其余维度大小一样，可以正常进行拼接；
如果基于y2的第一维度进行拼接，则会报错，因为子成员的其余维度大小不一样。

torch.stack[按创建新维度拼接]

touch.stack() 在创建新维度进行拼接，新创建的维度大小为拼接成员的数量，其中，所有成员的维度大小必须一样，如：

z3 = torch.stack(y1, dim=0) # 在第1维度新增维度，维度的大小为y1的数量
print(z3.shape)

out:
torch.Size([2, 8, 8, 3, 7, 7])

y1的子成员为2，在第一维度（dim=0）新增维度, 大小为2；
如果把y2进行拼接则会报错，因为y2的子成员的维度大小不一致。