Layer wise count of tensor using torch.count_nonzero()

vision
1

Lets say “a=torch.rand(1,4,2,2)” and “count_a=torch.zeros(4)”.

I want count of non zero of a(0,0,:,: ) to store in count_a(0), a(0,1,:,: ) to store in count_a(1) etc

If i used below code it stores total non zero count of “a” in to “count_a”.

Code:

import torch
count_a=torch.ones(4)
a=torch.rand(1,4,2,2)
count_a=torch.count_nonzero(a)
print("a value is ",a)
print("count is ",count_a)

Output:

a value is tensor([[[[0.8222, 0.4329],
[0.4672, 0.2243]],

     [[0.3086, 0.6668],
      [0.2699, 0.3848]],

     [[0.0251, 0.9675],
      [0.7624, 0.8663]],

     [[0.4905, 0.8086],
      [0.0784, 0.4926]]]])

count is tensor(16)

I want output to be tensor[4 4 4 4] as each layer has 4 non zero value instead of tesnor(16).

2

I don’t quite understand the expected output shape.
Your input tensor a has 16 values (1*4*2*2=16), so count_nonzero can return a value in [0, 16], while the expected output shape would contain more values.
Could you explain your use case a bit more?

3

If “a=torch.rand(1,4,2,2)” then think of 4 layer having 2x2 value. If “a=torch.rand(1,8,2,2)” then think 8 layer having 2x2 value.

I want layer wise non zero count value. so for “a=torch.rand(1,4,2,2)” if all 4 layer 2x2 value is non zero then i want output to be “count_a=[4 4 4 4]”. Sorry i forgot to mention that count_a is of size “torch.Size([4])” and i want output to be tensor([4., 4., 4., 4.])

If “a=torch.rand(1,8,2,2)”, if all the 8 layers 2x2 values are non zero then i want output to be tensor([4., 4., 4., 4., 4., 4., 4., 4. ]).

4

Thanks for the clarification!
I misunderstood the [4, 4, 4, 4] output as the shapes, not the values.
In that case, you could flatten the last two dimensions and use the dim argument in count_nonzero:

a = torch.rand(1,4,2,2)
count_a = torch.count_nonzero(a.view(1, 4, -1), 2)
print("a value is ",a)
print("count is ",count_a)
> count is  tensor([[4, 4, 4, 4]])
5

Thanks. That’s exactly i want.