PyTorch学习笔记: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
点个赞支持一下吧