在pytorch中沿轴计算非零元素

Question

我有一个(32x750)张量

tensor([[ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0043],
        [ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0043],
        [ 0.0000,  0.0044,  0.0000,  ...,  0.0044,  0.0000,  0.0000],
        ...,
        [ 0.0059,  0.0000,  0.0059,  ...,  0.0059,  0.0000,  0.0000],
        [ 0.0059,  0.0000,  0.0059,  ...,  0.0059,  0.0000,  0.0000],
        [ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0056,  0.0000]], device='cuda:0')

我想获取每行中非零元素的数量。类似的东西 [12 47 0 5 .... 8 7 50] 这个 discussion and 没有解决我的问题，关注 1-D tensor.

的非零元素的数量

谢谢

Answer 1

list_of_num_nonzero_in_each_row = []
for row in my tensor:
    list_of_num_nonzero_in_each_row.append(sum(row == 0.0).item())

Answer 2

使用这个 post 解决了问题我用过：750 - (tensor == 0).sum(dim=1)

Answer 3

Torch 现在具有 count_nonzero 功能 built in。这种矢量化实现将比迭代张量更快。它还支持对给定维度进行计数。

> torch.count_nonzero(x, dim=0)

在pytorch中沿轴计算非零元素

Count nonzeros element along an axis in pytorch

count

tensor