How to do masked mean while reserving the first dimension

vision
1

I have a tensor of size BxCxHxW. In the tensor there are some valid values which is larger than 0 and invalid values which equals to 0. Now I want to compute the mean of valid values in dimension CxHxW. So the result size is Bx1, the ith element equals the mean of valid values in the ith CxHxW tensor.

For example if I ignore C, I have a 2x1x3 tensor [[[0,1,2]], [[5,7,0]]], which means B=2, H=1, W=3. The desired result is a 2x1 tensor [1.5, 6], in which 1.5 is the average of 1 and 2, 6 is the average of 5 and 7. If I directly use tensor.mean(-1).mean(-1) it will count zeros into the result, which is not what I want.

2

a=torch.tensor([[[0,1,2]], [[5,7,0]]]).float()
mask=(a>0).float()
mask_mean = a.sum(-1) / mask.sum(-1)
output:
mask_mean=tensor([[1.5000], [6.0000]])
mask_mean.shape=torch.Size([2, 1])

2 Likes
3

One simple solution is as follows:

>>> a = torch.randn(10)
>>> a
tensor([-0.7635,  0.1626, -0.2671,  0.6118,  0.0304,  1.3959, -1.0044, -0.7236,
        -1.1573, -0.4817]
>>> aw = a > 0
>>> aw
tensor([0, 1, 0, 1, 1, 1, 0, 0, 0, 0], dtype=torch.uint8)
>>> torch.sum(a * aw.float()) / torch.sum(aw)
tensor(0.5502)
2 Likes
5

if it may happen, that you mask out everything (e.g. if you have multiple batches and some batchentries are completly masked out), you might want to do it like this:

batchsize = 32
datasize = 2
a = torch.randn(batchsize, datasize)
mask = (a > 0).float()
mask_sum = torch.sum(mask, dim=1)
mask_sum_modified = torch.clamp(mask_sum, min=1.0)
torch.sum(a * mask, dim=1) / mask_sum_modified