torch.cat 与torch.stack的区别

wenjurongyu2024-01-14 16:55

torch.cat 与torch.stack的区别

torch.cat

定义：按照规定的维度进行拼接。

实际使用: 例如使用BiLSTM时，将两个方向的向量进行叠加，就是用torch.cat。

python 复制代码

import torch

forward_lstm = torch.randn((2, 10, 3))
backward_lstm = torch.randn((2, 10, 3))

lstm_emd = torch.cat((forward_lstm, backward_lstm), dim=-1)

print(lstm_emd.size())
'''
torch.Size([2, 10, 6])
'''

torch.stack

定义：官方解释是在新的dim上进行叠加。叠加的意思就是增加一个维度。

本质：对张量进行unsqueeze(dim)之后，再进行torch.cat(dim=dim)操作。

实际使用：将张量合在一起，形成一个batch。

python 复制代码

import torch

batch_1 = torch.randn((10, 3))
batch_2 = torch.randn((10, 3))
batch = torch.stack((batch_1, batch_2), dim=0) 
print(batch.size()) 
'''
torch.Size([2, 10, 3])
'''

使用torch.unsqueeze 和torch.cat实现torch.stack功能

python 复制代码

import torch

batch_1 = torch.randn((10, 3))
batch_2 = torch.randn((10, 3))

batch_1 = torch.unsuqeeze(batch_1, dim=0)
batch_2 = torch.unsuqeeze(batch_2, dim=0)
batch = torch.cat((batch_1, batch_2), dim=0) 
print(batch.size()) 
'''
torch.Size([2, 10, 3])
'''