pytorch 合并与分割

cat 在原本的维度上扩展自己

import torch
import numpy as np
a = torch.rand(4,3,8)
b = torch.rand(5,3,8)
# dim 代表维度
print("合并第一维度后的 shape: {}".format(torch.cat([a,b],dim=0).shape))

合并第一维度后的 shape: torch.Size([9, 3, 8])

stack 会创建一个新的维度，要求合并的张量大小需要一致

import torch
import numpy as np
a = torch.rand(4,3)
b = torch.rand(4,3)
# dim 代表维度
print("合并后在最前面插入一个维度 shape: {}".format(torch.stack([a,b],dim=0).shape))
print("合并后在最后面插入一个维度 shape: {}".format(torch.stack([a,b],dim=-1).shape))

合并后在最前面插入一个维度 shape: torch.Size([2, 4, 3])
合并后在最后面插入一个维度 shape: torch.Size([4, 3, 2])

split

import torch
import numpy as np
a = torch.rand(4,3,8)
#按不同组合大小比例拆分 为 第一维度 4 拆分为 1 和 3
b , d = a.split([1,3],dim=0)
print(b.shape,d.shape)
#将第一维度 按单位数量为2 进行拆分
b , d = a.split(2,dim=0)
print(b.shape,d.shape)

torch.Size([1, 3, 8]) torch.Size([3, 3, 8])
torch.Size([2, 3, 8]) torch.Size([2, 3, 8])

chunk 按数量拆分

import torch
import numpy as np
a = torch.rand(4,3,8)

#将第一维度分成4份
b , d , e, f = a.chunk(4,dim=0)
print(b.shape,d.shape,e.shape,d.shape)

torch.Size([1, 3, 8]) torch.Size([1, 3, 8]) torch.Size([1, 3, 8]) torch.Size([1, 3, 8])