PyTorch学习笔记（四）：Tensor合并与分割

cat

a = torch.rand(4,32,8)
b = torch.rand(5,32,8)
print(torch.cat([a, b], dim=0).shape)    # torch.Size([9, 32, 8])

stack: create new dim

a = torch.rand(32,8)
b = torch.rand(32,8)
print(torch.stack([a, b], dim=0).shape)    # torch.Size([2, 32, 8])

split: by len

a = torch.rand(32,8)
b = torch.rand(32,8)
c = torch.stack([a, b], dim=0)
print(c.shape)      # torch.Size([2, 32, 8])

aa, bb = c.split([1, 1], dim=0)
print(aa.shape)     # torch.Size([1, 32, 8])
print(bb.shape)     # torch.Size([1, 32, 8])

aa_1, bb_1 = c.split([10, 22], dim=1)
print(aa_1.shape)    # torch.Size([2, 10, 8])
print(bb_1.shape)    # torch.Size([2, 22, 8])

chunk: by num

a = torch.rand(32,8)
b = torch.rand(32,8)
c = torch.stack([a,b], dim=0)
print(c.shape)       # torch.Size([2, 32, 8])

aa, bb = c.chunk(2, dim=0)
print(aa.shape)      # torch.Size([1, 32, 8])
print(bb.shape)      # torch.Size([1, 32, 8])

aa_1, bb_1 = c.chunk(2, dim=2)
print(aa_1.shape)    # torch.Size([2, 32, 4])
print(bb_1.shape)    # torch.Size([2, 32, 4])