如何计算 PyTorch 张量第二维中 1 和 0 的数量？

Question

我有一个大小为：的张量torch.Size([64, 2941])，它是 64 个批次，共 2941 个元素。

在所有 64 个批次中，我想计算张量第二维中 1 和 0 的总数，一直到第 2941 个，以便我将这些计数作为大小的张量torch.Size([2941])

我怎么做？

Answer 1

你可以将它们相加：

import torch
torch.manual_seed(2020)

# x is a fake data, it should be your tensor with 64x2941
x = (torch.rand((3,4)) > 0.5).to(torch.int32)
print(x)
# tensor([[0, 0, 1, 0],
#         [0, 0, 1, 1],
#         [0, 1, 1, 1]], dtype=torch.int32)

ones = (x == 1.).sum(dim=0)
print(ones)
# tensor([0, 1, 3, 2])

如果x是二进制，则可以通过简单的减法得到零的数量：

zeros = x.shape[1] - ones