torch.cat和torch.stack的区别

torch.cat和torch.stack的区别

stack操作后会在原来的基础上再增加一维，比如原来两个tensor的维度都是3维，经过stack后的结果为4维tensor； 而cat并不会改变维度数，只会在某一个维度上叠加。

import torch

# torch.cat函数
x = torch.rand((3, 240, 120))   
y = torch.rand((3, 240, 120))   
print(torch.cat((x, y), dim=0).shape)   # torch.Size([6, 240, 120])
print(torch.cat((x, y), dim=1).shape)   # torch.Size([3, 480, 120])
print(torch.cat((x, y), dim=2).shape)   # torch.Size([3, 240, 240])

# torch.stack函数
print(torch.stack((x, y), dim=0).shape)  # torch.Size([2, 3, 240, 120])
print(torch.stack((x, y), dim=1).shape)  # torch.Size([3, 2, 240, 120])
print(torch.stack((x, y), dim=2).shape)  # torch.Size([3, 240, 2, 120])