PyTorch 中如何使用 stack 方法？-JavaScript中文网-JavaScript教程资源分享门户

推荐答案

在 PyTorch 中，torch.stack 方法用于将一系列张量沿着一个新的维度进行堆叠。所有输入张量必须具有相同的形状。

// www.javascriptcn.com code example
import torch

# 创建两个形状相同的张量
tensor1 = torch.tensor([1, 2, 3])
tensor2 = torch.tensor([4, 5, 6])

# 沿着新维度（dim=0）堆叠张量
stacked_tensor = torch.stack((tensor1, tensor2), dim=0)

print(stacked_tensor)

输出结果：

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

本题详细解读

1. `torch.stack` 的基本用法

torch.stack 方法的基本语法如下：

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

tensors：一个包含多个张量的序列（如列表或元组），这些张量必须具有相同的形状。
dim：指定新维度的位置，沿着该维度进行堆叠。dim 可以是负数，表示从最后一个维度开始计数。
out：可选参数，指定输出张量。

2. `torch.stack` 与 `torch.cat` 的区别

torch.stack 会创建一个新的维度，并将输入张量沿着该维度堆叠。
torch.cat 则是在现有维度上连接张量，不会创建新的维度。

3. 示例代码解析

// www.javascriptcn.com code example
import torch

# 创建两个形状相同的张量
tensor1 = torch.tensor([1, 2, 3])
tensor2 = torch.tensor([4, 5, 6])

# 沿着新维度（dim=0）堆叠张量
stacked_tensor = torch.stack((tensor1, tensor2), dim=0)

print(stacked_tensor)

在这个例子中，tensor1 和 tensor2 都是形状为 (3,) 的一维张量。通过 torch.stack 方法，我们将它们沿着 dim=0 堆叠，生成了一个形状为 (2, 3) 的二维张量。

4. 不同维度的堆叠

你可以通过改变 dim 参数来指定不同的堆叠维度。例如：

# 沿着 dim=1 堆叠
stacked_tensor = torch.stack((tensor1, tensor2), dim=1)

print(stacked_tensor)

输出结果：

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

在这个例子中，dim=1 表示沿着第二个维度（即列）进行堆叠，生成了一个形状为 (3, 2) 的二维张量。

5. 注意事项

所有输入张量必须具有相同的形状。
dim 参数决定了新维度的位置，必须确保 dim 的值在合理范围内（即不超过输入张量的维度数加一）。

PyTorch 中如何使用 stack 方法？