PyTorch-获取3D张量中2D张量和的列表

Question

我有一个由2D张量组成的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]]
])

我需要一个2D张量和的列表或张量，例如：sums = [3, 3, 3]。到目前为止，我有：

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

我认为这只能用PyTorch来完成，但我已经尝试过使用具有所有可能维度的torch.sum()，并且我总是得到这些2D张量的各个字段的和，例如：

[[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()

输出：[3, 3, 3]