Как подсчитать количество единиц и нулей во втором измерении тензора PyTorch?

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