Pytorch基础 - 5. torch.cat() 和 torch.stack()

目录

1. torch.cat(tensors, dim)

2. torch.stack(tensors, dim)

3. 两者不同

---

torch.cat() 和 torch.stack()常用来进行张量的拼接，在神经网络里经常用到。且前段时间有一个面试官也问到了这个知识点，虽然内容很小很细，但需要了解。

1. torch.cat(tensors, dim)

• tensors：待拼接的多个张量，可用list, tuple表示
• dim：待拼接的维度，默认是0
• 注意：tensors里不同张量对应的待拼接维度的size可以不一致，但是其他维度的size要保持一致。如代码中待拼接维度是0，x和y对应的维度0上的值不一样，但是其他维度上的值(维度1上的值)要保持一致，即都为4，否则会报错。

示例：新生成的tensor在dim=0这个维度进行了拼接，即 3 + 2 = 5，剩余维度保持不变

x = torch.rand(3, 4)
y = torch.rand(2, 4)
xy = torch.cat([x, y], dim=0)   
print(xy.shape)   # torch.Size([5, 4])

2. torch.stack(tensors, dim)