一、问题引出
好久没更新啦!最近在学习沐神《动手学深度学习》6.5节池化层的时候,发现沐神在两处相似的地方使用了两种Python拼接函数torch.cat()和torch.stack():
百思不得其解,于是查阅相关文档之后终于弄清楚了两者之间的区别,遂做总结如下。
二、问题解决
1.torch.cat()
torch.cat()函数可以将多个张量拼接成一个张量。torch.cat()有两个参数,第一个是要拼接的张量的列表或是元组;第二个参数是拼接的维度
python
# 假设是时间步T1的输出
T1 = torch.tensor([[1, 2, 3],
[4, 5, 6],
[7, 8, 9]])
# 假设是时间步T2的输出
T2 = torch.tensor([[10, 20, 30],
[40, 50, 60],
[70, 80, 90]])
print("T1.shape: ", T1.shape, "T2.shape: ", T2.shape)
print(torch.cat((T1,T2),dim=0).shape)
print(torch.cat((T1,T2),dim=1).shape)
输出为:
2.torch.stack()
torch.stack()函数同样有张量列表和维度两个参数。stack与cat的区别在于,torch.stack()函数要求输入张量的大小完全相同,得到的张量的维度会比输入的张量的大小多1,并且多出的那个维度就是拼接的维度,那个维度的大小就是输入张量的个数。
python
print("T1.shape: ", T1.shape, "T2.shape: ", T2.shape)
print(torch.stack((T1,T2),dim=0).shape)
print(torch.stack((T1,T2),dim=1).shape)
print(torch.stack((T1,T2),dim=2).shape)