torch.cat()与torch.stack()——数组的拼接

torch.cat()

torch.cat(tensors, dim=0, *, out=None) → Tensor

官方解释：利用给定的维度连接给定的数组序列(cat代表concatenate)，所有数组必须具有相同的形状（连接维度除外）或为空。
相当于按指定维度将数组进行拼接

参数解释：

• tensors：要连接的数组序列(元组tuple或者列表list)
• dim：数组连接的维度
• out：输出数组(一般用不到，如果有输出，则可以直接进行赋值操作)

注意：
①tensors输入的必须是数组序列，不能是单个数组；
②输入的数组序列除了dim维度，其他维度必须形状相同。

举例：

import torch
a=torch.arange(6).reshape(2,3)
b=torch.arange(12)
c=torch.cat((a,b.reshape(4,3)),dim=0)
# 沿第0维度进行拼接，也就是按行拼接(竖着拼)
d=torch.cat((a,b.reshape(2,6)),dim=1)
# 沿第1维度进行拼接，也就是按列拼接(横着拼)
print(c)
print(c.shape)
print(d)
print(d.shape)

输出：

tensor([[ 0,  1,  2],
        [ 3,  4,  5],
        [ 0,  1,  2],
        [ 3,  4,  5],
        [ 6,  7,  8],
        [ 9, 10, 11]])
torch.Size([6, 3])
tensor([[ 0,  1,  2,  0,  1,  2,  3,  4,  5],
        [ 3,  4,  5,  6,  7,  8,  9, 10, 11]])
torch.Size([2, 9])

利用torch.cat()沿dim拼接，在形状上看相当于对dim进行相加，其余维度大小不变，利用这个思想，可以很容易理解高维数组的拼接

高维举例：

import torch
a=torch.ones(4*256*56*56).reshape(4,256,56,56)
b=torch.arange(4*128*56*56).reshape(4,128,56,56)
c=torch.zeros(4*64*56*56).reshape(4,64,56,56)
d=torch.cat((a,b,c),dim=1)
print(d.shape)

输出：

torch.Size([4, 448, 56, 56])

上述例子在卷积神经网络中常用于特征图的堆叠。

torch.stack()

torch.stack(tensors, dim=0, *, out=None) → Tensor

官方解释：沿着新的维度连接一系列数组，所有的数组都需要具有相同的大小。
相当于先将多个n维数组进行扩维操作，然后再拼接为一个n+1维的数组

参数解释：

• tensors：要连接的数组序列(元组tuple或者列表list)
• dim：要插入的维度，大小必须介于0和需要拼接的数组维数之间(dim最大不超过数组的维数)
• out：输出数组(与cat()类似，一般用不到)

注意：
①与cat类似，必须输入数组序列，不能是单个数组；
②输入的所有数组序列形状(尺寸)必须一致(这里与cat有区别)。

举例：

import torch
a=torch.arange(12).reshape(3,4)
b=torch.ones(12).reshape(3,4)
c=torch.stack((a,b),dim=0)
d=torch.stack((a,b),dim=1)
e=torch.stack((a,b),dim=2)
# dim最大可到输入数组的维数，即a、b的维数
print(c)
print(c.shape)
print(d)
print(d.shape)
print(e)
print(e.shape)

输出：

tensor([[[ 0.,  1.,  2.,  3.],
         [ 4.,  5.,  6.,  7.],
         [ 8.,  9., 10., 11.]],

        [[ 1.,  1.,  1.,  1.],
         [ 1.,  1.,  1.,  1.],
         [ 1.,  1.,  1.,  1.]]])
torch.Size([2, 3, 4])
tensor([[[ 0.,  1.,  2.,  3.],
         [ 1.,  1.,  1.,  1.]],

        [[ 4.,  5.,  6.,  7.],
         [ 1.,  1.,  1.,  1.]],

        [[ 8.,  9., 10., 11.],
         [ 1.,  1.,  1.,  1.]]])
torch.Size([3, 2, 4])
tensor([[[ 0.,  1.],
         [ 1.,  1.],
         [ 2.,  1.],
         [ 3.,  1.]],

        [[ 4.,  1.],
         [ 5.,  1.],
         [ 6.,  1.],
         [ 7.,  1.]],

        [[ 8.,  1.],
         [ 9.,  1.],
         [10.,  1.],
         [11.,  1.]]])
torch.Size([3, 4, 2])

仔细观察上个案例维度的变化，可以发现当输入为两组数组时，dim定为几，拼接后哪个维度就是2(有两个输入数组)，相当于做了一个扩维拼接操作。首先按dim增加一个维度，然后再从该维度上进行拼接操作。

cat与stack的区别

torch.cat()是直接在原数组数据上进行拼接，不会改变维数大小；torch.stack首先进行扩维，然后再进行拼接，会将维数增大一个单位。

官方文档

torch.cat():https://pytorch.org/docs/stable/generated/torch.cat.html
torch.stack():https://pytorch.org/docs/stable/generated/torch.stack.html

点个赞支持一下吧