python - 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]]
                  ])

我需要这些二维张量之和的列表或张量，例如：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 中做到这一点？

标签： pythonpytorch

Answer 1

您可以通过将 dims 作为元组传递来立即执行此操作。

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

或列表

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