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

时间:2020-08-10 07:03:17

标签: python pytorch

我有一个由2D张量组成的3D张量,例如。 g。:

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

我需要这些二维张量的列表或总和张量,例如g.:sums = [3, 3, 3]。到目前为止,我有:

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

我认为这只能通过PyTorch来完成,但是我尝试使用torch.sum()并具有所有可能的尺寸,而且我总是在这些2D张量的各个字段上获得总和,例如。 g。:

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

如何在PyTorch中完成此操作?

2 个答案:

答案 0 :(得分:3)

您可以通过将昏暗作为元组传递一次。

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

或用于列表

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

答案 1 :(得分:1)

如果正确理解了您的问题,这应该可以完成工作:

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

输出: [3, 3, 3]