PyTorch的torch.stack 和 torch.cat

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

  • torch.stack 会沿着一个新的维度将输入张量堆叠起来。所有的输入张量必须具有相同的形状。
  • dim 参数指定了新张量增加的维度索引。默认是0,意味着新维度会被添加到张量的最前面。
  • 输出张量的维度会增加一个,形状会在指定的维度上增加一个与输入张量数量相同的维度。

有两个形状为 [a, b] 的张量,如果使用 torch.stack 来堆叠它们,那么输出张量的形状将是 [2, a, b]

torch.cat(tensors, dim=0, out=None):

  • torch.cat 会沿着指定的维度将输入张量连接起来。所有输入张量在除了连接维度外的其他维度上必须具有相同的形状。
  • dim 参数指定了沿着哪个维度进行连接。
  • 输出张量的形状与输入张量在连接维度上的形状之和相同,其他维度保持不变。

有两个形状为 [a, b] 的张量,如果使用 torch.cat 在维度0上进行连接,那么输出张量的形状将是 [2*a, b]

你可能感兴趣的:(pytorch,机器学习)