pytorch中常见的合并和分割

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()