import torch
'''
1-数据合并:cat(沿着维度);stack(增加维度)
2-数据分割:split(按照长度);chunk(按照数量)
'''
a1 = torch.rand(2, 3, 28, 28)
a2 = torch.rand(4, 3, 28, 28)
'''
torch.cat([a1, a2], dim=0)
dim = 指定拼接的维度
注意,使用torch.cat进行拼接时除了拼接维度可以不同外,其他的维度必须相同
'''
def zqb_cat():
print(a1.shape, a2.shape) # torch.Size([2, 3, 28, 28]) torch.Size([4, 3, 28, 28])
a3 = torch.cat([a1, a2], dim=0)
print(a3.shape) # torch.Size([6, 3, 28, 28])
# a4 = torch.cat([a1, a2],dim=1) #RuntimeError: Sizes of tensors must match except in dimension 1.
'''
stack
沿新维度连接一系列张量
所有张量的size必须一样
使用stack操作后,会在传入的dim前面插入一个新的维度,而其余原本维度的形状不变
'''
def zqb_stack():
a1 = torch.rand(4, 3, 28, 28)
a2 = torch.rand(4, 3, 28, 28)
a3 = torch.stack([a1, a2], dim=0) # torch.Size([2, 4, 3, 28, 28])
print(a3.shape)
a4 = torch.stack([a1, a2], dim=1)
print(a4.shape) # torch.Size([4, 2, 3, 28, 28])
'''
torch.split(a2, [1, 2, 1], dim=0)
按照长度进行分割,指定每一个分割的长度,和沿着那个维度分割
'''
def zqb_split():
print(a2.shape) # torch.Size([4, 3, 28, 28])
a3, a4, a5 = torch.split(a2, [1, 2, 1], dim=0)
print(a3.shape, a4.shape,
a5.shape) # torch.Size([1, 3, 28, 28]) torch.Size([2, 3, 28, 28]) torch.Size([1, 3, 28, 28])
'''
torch.chunk(a2,2,dim=0)
按照数量进行分割,指定分割的数量和分割的维度
'''
def zqb_chunk():
print(a2.shape) # torch.Size([4, 3, 28, 28])
a3, a4 = torch.chunk(a2, 2, dim=0)
print(a3.shape, a4.shape) # ([2, 3, 28, 28]) torch.Size([2, 3, 28, 28])
if __name__ == "__main__":
zqb_cat()
zqb_stack()
zqb_split()
zqb_chunk()
pytorch中常见的合并和分割
最新推荐文章于 2024-07-18 20:02:41 发布