Torch Concat详解

AEGJO • 2025-04-12 13:00 • 编程

一、拼接张量

拼接（Concatenation）张量是将两个张量沿着某个维度进行拼接，得到一个更大的张量。在PyTorch中，可以使用torch.cat来完成拼接张量的操作。

import torch

# 创建3 x 2的张量
x = torch.randn(3, 2)

# 创建3 x 3的张量
y = torch.randn(3, 3)

# 沿着第二个维度对两个张量进行拼接
z = torch.cat([x, y], dim=1)

print(z)

在上面的例子中，我们先使用torch.randn创建了两个不同的张量x和y，张量x的维度是3 x 2，张量y的维度是3 x 3。使用torch.cat将张量x和y沿着第二个维度（即列）拼接，得到了一个维度为3 x 5的新张量z。

二、注意事项

在使用torch.cat进行张量拼接时，需要注意以下几点。

拼接的维度的大小必须相同，除拼接维度外，其他维度大小也必须相同。
拼接的维度的编号必须在0到张量维度数减1的范围内。
拼接的维度大小可以根据需要设置为-1，此时大小将自动推断。
如果两个张量是CPU张量，则拼接后的张量也是CPU张量。如果两个张量是CUDA张量，则拼接后的张量也是CUDA张量。

三、拼接多个张量

我们也可以使用torch.cat来拼接多个张量。下面的例子将展示如何同时拼接三个张量。

import torch

# 创建3 x 2的张量
x = torch.randn(3, 2)

# 创建3 x 3的张量
y = torch.randn(3, 3)

# 创建3 x 4的张量
z = torch.randn(3, 4)

# 沿着第二个维度对三个张量进行拼接
w = torch.cat([x, y, z], dim=1)

print(w)

在上面的例子中，我们分别创建了3个不同大小的张量，使用torch.cat将它们沿着第二个维度（即列）拼接成一个维度为3 x 9的张量w。

四、使用stack拼接张量

如果需要在新创建的维度上拼接张量，可以使用torch.stack。栈（Stack）张量是一个新的张量，它将输入张量沿着新创建的维度进行堆叠。

import torch

# 创建3 x 2的张量
x = torch.randn(3, 2)

# 创建3 x 2的张量
y = torch.randn(3, 2)

# 沿着新维度将两个张量进行堆叠
z = torch.stack([x, y], dim=0)

print(z)

在上面的例子中，我们先使用torch.randn创建了两个不同的张量x和y，张量x和张量y的维度都是3 x 2。使用torch.stack将张量x和张量y沿着新维度（即第0个维度）堆叠，得到了一个维度为2 x 3 x 2的新张量z。

五、结论

在PyTorch中，torch.cat和torch.stack是非常有用的函数，它们可以方便地对多个张量进行拼接操作。在使用这两个函数时需要注意维度的大小和编号，以及张量的类型。

原创文章，作者：AEGJO，如若转载，请注明出处：https://www.506064.com/n/369242.html