拆分
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的子成员的维度大小不一致。