1
$\begingroup$

I have a tensor of shape = [a,b,c]

The tensor mean along dim=1 will give me an output of shape = [a,c]

My goal is to compute the mean of values along dim=1 greater than a threshold. How is this possible in TensorFlow?

I am aware of tf.where() and tf.boolean_mask() functions but not sure exactly how to use them for this task.

$\endgroup$

1 Answer 1

0
$\begingroup$

First, If you calculate the mean along dim=1 the output shape should be [a, c].

If you want to mask the mean that's less then a threshold and set it to zero you can do

# generate data
torch.manual_seed(42)
a, b, c = 2, 3, 4
t = torch.normal(torch.ones(a, b, c))
print(t)

tensor([[[ 2.9269,  2.4873,  1.9007, -1.1055],
         [ 1.6784, -0.2345,  0.9569, -0.6047],
         [ 1.3559,  0.3134,  0.5066,  1.2415]],

        [[-0.1109,  1.0915, -1.3169,  0.7832],
         [ 0.6903,  0.6043,  1.8034,  0.3784],
         [ 0.4080,  0.9369,  0.1714,  1.3309]]])

to get the mean and mask to 0

mean = t.mean(dim=1)
print(mean)

tensor([[ 1.9871,  0.8554,  1.1214, -0.1562],
        [ 0.3291,  0.8776,  0.2193,  0.8308]])

threshold = 1
mean[mean<threshold] = 0
print(mean)

tensor([[1.9871, 0.0000, 1.1214, 0.0000],
        [0.0000, 0.0000, 0.0000, 0.0000]])

If you instead just want a list of means you can do this:

print(mean[mean<threshold])

tensor([1.9871, 1.1214])
$\endgroup$
1
  • $\begingroup$ I am asking solution for TF not pytorch. Nevertheless my goal is not to mask means above a threshold but to compute mean for values above a threshold. $\endgroup$
    – Astro
    Commented May 6, 2020 at 9:58

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.