PyTorch常用函数torch.cat()中dim参数使用说明

Part 1： 简介

在PyTorch中，torch.cat()是一个被广泛使用的函数。它可以让我们在某个维度上把多个张量组合在一起。对于那些想要深入了解使用PyTorch进行数据分析和建模的开发者来说，理解torch.cat()函数的dim参数是非常重要的。

在PyTorch中，几乎所有与神经网络有关的操作都涉及到张量（Tensor）操作。因此，在PyTorch中，将多个相同形状的张量沿某个轴/维度连接起来的过程非常重要。这就是 torch.cat() 函数的作用。torch.cat() 的最基本用法如下：

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

其中tensors表示要拼接的张量列表，dim表示我们希望在哪个维度上连接，默认是0，即在第一维上连接。out是输出张量，可不传入，当传入此参数时其大小必须能容纳在cat操作后的输出tensor中。

Part 2： dim参数的说明

dim参数指示拼接发生的轴或维度。在拼接多个张量时，我们必须指定在哪个维度上拼接它们。dim参数可以是正数、负数或None（默认为0），具体来说，dim参数可以有以下三种常见用法：

正数

最常见的方式是使用正整数来指定要连接的维度/轴的索引值。例如，在将两个大小为 3x5x7 的张量沿第2个维度拼接在一起时，这些张量变成一个形状为 3x10x7 的张量。

# 定义两个大小都为[3, 5, 7]的随机Tensor
tensor1 = torch.randn(3, 5, 7)
tensor2 = torch.randn(3, 5, 7)
# 在第二维度上（索引1）进行合并
cat_tensor = torch.cat((tensor1, tensor2), dim=1)
print(cat_tensor.shape) # 输出: torch.Size([3, 10, 7])

负数

我们也可以使用负整数来表示要连接的轴/维度。当dim参数被设置为负整数时，它代表距离张量最后一个轴的间隔数。例如，将一个大小为3x5x7 和一个大小为3x6x7的张量沿着最后一个维度进行拼接，即 concatenate 第三个维度：

# 定义两个大小分别为 [3, 5, 7], [3, 6, 7] 的随机Tensor
tensor1 = torch.randn(3, 5, 7)
tensor2 = torch.randn(3, 6, 7)
# 在最后一个维度上（-1表示）进行合并
cat_tensor = torch.cat((tensor1, tensor2), dim=-1)
print(cat_tensor.shape) # 输出: torch.Size([3, 5, 14])

None

如果 dim 参数的值为 None，则会将所有输入张量沿着前面的维度全部展开。这通常会在神经网络模型中使用，例如在线性层之间堆叠各个特征向量时。

# 定义两个大小分别为 [3, 5, 7], [4, 6, 8] 的随机Tensor
tensor1 = torch.randn(3, 5, 7)
tensor2 = torch.randn(4, 6, 8)
# 将每个张量reshape为1D向量
resized_t1 = tensor1.view(-1)
resized_t2 = tensor2.view(-1)
# 按行连接两个1D张量  
cat_tensor = torch.cat((resized_t1, resized_t2), dim=None)
print(cat_tensor.shape) # 输出: torch.Size([315])

Part 3： 总结

torch.cat() 函数是PyTorch非常有用的函数之一，它可以在某个维度上将多个张量组合成一个大张量。理解dim参数的含义和使用方法对于深入学习PyTorch和构建神经网络非常重要。通过在 dim 参数上增加或减少索引来改变连接选定的张量的方式，我们可以让torch.cat()函数在数据处理、模型设计和深度学习中发挥重要作用。

来源：https://juejin.cn/post/7222897518501150775