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

Question

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

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

我该怎么做？

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

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

How can I count the number of 1's and 0's in the second dimension of a PyTorch Tensor?

python

pytorch

tensor