python - 如何计算 PyTorch 张量的第二维中 1's and 0' 的数量？

Question

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

在所有 64 个批处理中，我想计算张量第二维中 1 和 0 的总数，一直到第 2941 个，以便我将这些计数作为大小为 torch 的张量。尺寸([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