Pytorch torch.stack()和torch.cat()函数的区别

在本文中，我们将介绍Pytorch中的两个常用函数torch.stack()和torch.cat()的区别以及使用示例。

torch.stack()

torch.stack()函数用于沿着新的维度对输入的张量序列进行堆叠。具体而言，它将一系列张量按照指定的维度进行连接，并创建一个新的张量。

语法如下：

torch.stack(tensors, dim=0, *, out=None)

其中，参数说明如下：
– tensors: 张量序列，可以是一个列表、元组或迭代器。
– dim: 指定沿着哪个维度进行堆叠，默认为0。

下面是一个使用torch.stack()函数的示例：

import torch

# 创建两个张量
a = torch.tensor([1, 2, 3])
b = torch.tensor([4, 5, 6])

# 使用torch.stack()函数进行堆叠
c = torch.stack((a, b))
print(c)

输出结果为：

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

在示例中，我们首先创建了两个一维张量a和b。然后使用torch.stack()函数对它们进行了堆叠，得到了一个二维张量c。可以看到，新的张量c的第一个维度是2，对应于堆叠后的张量个数，而第二个维度与原始张量的维度保持一致。

torch.cat()

torch.cat()函数用于沿着现有维度对输入的张量序列进行连接。具体而言，它将一系列张量按照指定的维度进行连接，并创建一个新的张量。

语法如下：

torch.cat(tensors, dim=0, out=None)

其中，参数说明与torch.stack()函数一致。

下面是一个使用torch.cat()函数的示例：

import torch

# 创建两个张量
a = torch.tensor([1, 2, 3])
b = torch.tensor([4, 5, 6])

# 使用torch.cat()函数进行连接
c = torch.cat((a, b))
print(c)

输出结果为：

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

在示例中，我们同样创建了两个一维张量a和b。然后使用torch.cat()函数对它们进行了连接，得到了一个一维张量c。可以看到，新的张量c包含了原始张量a和b中的所有元素。

区别分析

torch.stack()和torch.cat()函数的最主要区别在于它们进行连接的方式。

torch.stack()在连接张量时会增加一个新的维度，将每个张量堆叠在新的维度上。因此，堆叠后的张量维度会增加。
torch.cat()在连接张量时不会增加新的维度，而是将张量按照指定的维度进行连接。因此，连接后的张量在连接维度上的大小会增加，而其他维度大小不变。

通过下面的示例来进一步理解这两个函数的区别：

import torch

# 创建两个二维张量
a = torch.tensor([[1, 2, 3],
                  [4, 5, 6]])
b = torch.tensor([[7, 8, 9],
                  [10, 11, 12]])

# 使用torch.stack()对二维张量进行堆叠
c = torch.stack((a, b), dim=0)
print(c)

# 使用torch.stack()对二维张量进行堆叠
d = torch.cat((a, b), dim=0)
print(d)

输出结果为：

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

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

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

在示例中，我们创建了两个二维张量a和b。使用torch.stack()函数对二维张量进行堆叠时，新的张量c的维度增加了一维，堆叠后的张量为一个三维张量。而使用torch.cat()函数进行连接时，新的张量d的维度与原始张量的维度相同，只是在连接维度上的大小增加了。