为什么在维度（大小为 1）上使用“：”和“0”时会得到不同的输出？

Question

我认为最后两个语句的输出（使用

:

与

0

相比）会是相同的，因为只有一维。

x = torch.arange(1,10).reshape(3,1,3)
x[:,:,2], x[:,0,2]

(tensor([[3],
         [6],
         [9]]),
tensor([3, 6, 9]))

有人可以向我解释一下不同输出的原因吗？

Answer 1

您的起始张量是 3D 且形状为

(3,1,3)

。两个索引表达式的区别是：

• 在前一个（

  x[:,:,2]

  ）中，您选择

  dim=0

  和

  dim=1

  中的所有元素，但仅选择i=2

  上的

  dim=2

  。

• 而在后者 (

  x[:,0,2]

  ) 中，您选择

  dim=0

  中的所有元素，但仅选择

  i=0

  上的元素

  dim=1

  和

  i=2

  上的元素

  dim=2

  。由于

  dim=1

  上的选择是单例，因此与其他表达式相比，您实际上多了 挤压 一个维度。

因此，第一个表达式产生 2D 张量，而第二个表达式返回 1D 张量。