PyTorch - 獲取 3D 張量中二維張量總和的列表

Question

我有一個由二維張量組成的 3D 張量，例如：

t = torch.tensor([[[0, 0, 1],
                   [0, 1, 0],
                   [1, 0, 0]],

                  [[0, 0, 1],
                   [0, 1, 0],
                   [1, 0, 0]],

                  [[0, 0, 1],
                   [0, 1, 0],
                   [1, 0, 0]]
                  ])

我需要這些二維張量之和的列表或張量，例如： sums = [3, 3, 3] 。 到目前為止，我有：

sizes = [torch.sum(t[i]) for i in range(t.shape[0])]

我認為這可以僅使用 PyTorch 來完成，但我已經嘗試使用具有所有可能維度的torch.sum()並且我總是得到這些二維張量的各個字段的總和，例如：

[[0, 0, 3],
[0, 3, 0],
[3, 0, 0]]

如何在 PyTorch 中做到這一點？

Answer 1

您可以通過將 dims 作為元組傳遞來立即執行此操作。

t.sum(dim=(0,1))
tensor([3, 3, 3])

或列表

t.sum(dim=(0,1)).tolist()
[3, 3, 3]

Answer 2

如果正確理解您的問題，這應該可以完成工作：

t.sum(0).sum(1).tolist()

Output: [3, 3, 3]