torch.stack() 是 PyTorch 中的一个函数,用于将多个张量沿着新的维度进行堆叠。具体而言,torch.stack(seq, dim=0) 执行以下操作:
将 seq 序列中的每个张量视为一行,构建一个新的二维张量(即矩阵),其中第
i
i
i 行对应于第
i
i
i 个张量;
在指定的 dim 维度上对这些行进行堆叠,从而得到一个新的张量。
例如,假设有两个形状为 [3] 的张量 a 和 b,它们的数值如下所示:
import torch
a = torch.tensor([1, 2, 3])
b = torch.tensor([4, 5, 6])
那么,我们可以使用 torch.stack([a, b], dim=0) 函数将这两个张量在新的第一维度上进行堆叠,得到一个形状为 [2, 3] 的张量,如下所示:
output = torch.stack([a, b], dim=0)
print(output)
输出结果如下:
tensor([[1, 2, 3],
[4, 5, 6]])
其中,第一行表示 a 张量中的元素,第二行表示 b 张量中的元素。
需要注意的是,torch.stack() 函数要求所有被堆叠的张量的形状必须相同。此外,堆叠后的张量的维度数目将比原先多一维。