torch.stack()

顾昊穹
2023-12-01

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() 函数要求所有被堆叠的张量的形状必须相同。此外,堆叠后的张量的维度数目将比原先多一维。

 类似资料:

相关阅读

相关文章

相关问答