在 PyTorch 中,可以使用 torch.cattorch.stack 来合并多个张量(tensor)。两者的主要区别是合并方式不同:

  1. torch.cat:在给定维度上将张量连接起来,不会创建新的维度。
  2. torch.stack:在给定维度上增加一个新维度来堆叠张量。

示例代码

使用 torch.cat 合并张量
import torch

# 假设有两个张量
tensor1 = torch.tensor([[1, 2], [3, 4]])
tensor2 = torch.tensor([[5, 6], [7, 8]])

# 在第 0 维度(行)合并
result_cat_0 = torch.cat((tensor1, tensor2), dim=0)
print(result_cat_0)
# 输出:
# tensor([[1, 2],
#         [3, 4],
#         [5, 6],
#         [7, 8]])

# 在第 1 维度(列)合并
result_cat_1 = torch.cat((tensor1, tensor2), dim=1)
print(result_cat_1)
# 输出:
# tensor([[1, 2, 5, 6],
#         [3, 4, 7, 8]])

使用 torch.stack 堆叠张量

# 在第 0 维度增加新维度堆叠
result_stack_0 = torch.stack((tensor1, tensor2), dim=0)
print(result_stack_0)
# 输出:
# tensor([[[1, 2],
#          [3, 4]],
#         [[5, 6],
#          [7, 8]]])

# 在第 1 维度增加新维度堆叠
result_stack_1 = torch.stack((tensor1, tensor2), dim=1)
print(result_stack_1)
# 输出:
# tensor([[[1, 2],
#          [5, 6]],
#         [[3, 4],
#          [7, 8]]])

注意事项

  1. 使用 torch.cat 时,张量在非拼接维度上的大小必须相同。
  2. 使用 torch.stack 时,所有张量的形状必须完全相同。

选择方法取决于你的需求:

  • 如果需要简单地将张量扩展到一个现有维度,使用 torch.cat
  • 如果需要增加一个新维度来堆叠张量,使用 torch.stack
Logo

有“AI”的1024 = 2048,欢迎大家加入2048 AI社区

更多推荐