torch.cat 与torch.stack的区别

torch.cat 与torch.stack的区别

torch.stack的就是torch.unsqueeze+torch.cat操作。二者都要求拼接的张量形状相同。

torch.cat:形状是[a, b]的两个张量，dim=0:形成[a * 2, b] , dim=1：形成[a, b* 2]。

import torch

a = torch.randn((2, 3))
b = torch.randn((2, 3))
c = torch.cat((a, b), dim=0)
d = torch.cat((a, b), dim=1)
print(c.size(), d.size())
'''
torch.Size([4, 3]) torch.Size([2, 6])
'''

torch.stack，官方解释是在新的dim上进行拼接。本质是相当于张量unsqueeze(dim)之后，torch.cat(dim=dim)

import torch

a = torch.randn((2, 3))
b = torch.randn((2, 3))
c = torch.stack((a, b), dim=2)
d = torch.cat((a.unsqueeze(2), b.unsqueeze(2)), dim=2) # 和stack的操作结果是一致的。
print(c.size(), d.size())
print(c)
print(d)
'''
torch.Size([2, 3, 2]) torch.Size([2, 3, 2])
tensor([[[-0.5118,  0.4662],
         [-1.1902, -0.3203],
         [ 0.0080, -1.7615]],

        [[ 1.8801, -0.1554],
         [-0.0998,  1.0836],
         [ 0.9182, -1.0683]]])
tensor([[[-0.5118,  0.4662],
         [-1.1902, -0.3203],
         [ 0.0080, -1.7615]],

        [[ 1.8801, -0.1554],
         [-0.0998,  1.0836],
         [ 0.9182, -1.0683]]])
         '''