Pytorch torch.stack()和torch.cat()函数有什么区别

在本文中，我们将介绍Pytorch中的两个重要函数torch.stack()和torch.cat()的区别以及如何正确使用它们。

torch.stack()

torch.stack()是将一系列张量按照给定的维度进行拼接，在新维度上创建一个新的张量。通过torch.stack()函数，我们可以将多个张量按照指定的维度进行堆叠。它会创建一个新的张量，然后将输入张量堆叠到新的维度上。

import torch

tensor1 = torch.tensor([[1, 2, 3], [4, 5, 6]])
tensor2 = torch.tensor([[7, 8, 9], [10, 11, 12]])

stacked_tensor = torch.stack((tensor1, tensor2), dim=0)
print(stacked_tensor)

输出结果为：

tensor([[[ 1,  2,  3],
         [ 4,  5,  6]],

        [[ 7,  8,  9],
         [10, 11, 12]]])

在上面的示例中，我们有两个形状为(2, 3)的张量tensor1和tensor2。我们使用torch.stack()函数将这两个张量在新的维度0上进行堆叠，得到一个新的三维张量stacked_tensor。stacked_tensor的形状为(2, 2, 3)，其中第一个维度表示堆叠的张量个数。

torch.cat()

torch.cat()函数用于按照指定的维度连接张量。与torch.stack()不同，torch.cat()函数将输入张量按照给定的维度直接连接在一起，而不是创建一个新的维度。需要注意的是，torch.cat()函数要求其他维度的形状必须相同。

import torch

tensor1 = torch.tensor([[1, 2, 3], [4, 5, 6]])
tensor2 = torch.tensor([[7, 8, 9], [10, 11, 12]])

concatenated_tensor = torch.cat((tensor1, tensor2), dim=0)
print(concatenated_tensor)

输出结果为：

tensor([[ 1,  2,  3],
        [ 4,  5,  6],
        [ 7,  8,  9],
        [10, 11, 12]])

在上面的示例中，我们同样有两个形状为(2, 3)的张量tensor1和tensor2。我们使用torch.cat()函数将这两个张量按照维度0直接连接在一起，得到一个新的形状为(4, 3)的张量concatenated_tensor。

区别和适用场景

torch.stack()函数和torch.cat()函数的区别主要有以下几点：

维度要求不同：torch.stack()函数要求输入张量具有相同的维度，新张量会在一个新的维度上堆叠；而torch.cat()函数要求其他维度的形状必须相同，直接在指定维度上进行连接。
输出张量维度不同：torch.stack()函数会创建一个新的维度用于堆叠张量，输出张量的维度会增加；而torch.cat()函数不会创建新的维度，输出张量的维度不变。
适用场景不同：torch.stack()函数适用于需要在新的维度上堆叠多个张量的情况，常用于创建更高维度的张量；torch.cat()函数适用于在已有的维度上直接连接张量，常用于拼接同一维度上的张量。

根据以上区别，我们可以根据不同的需求选择使用适当的函数。