How to average only non-zero entries in Tensor?

user_9758922 0
Viewed 2694

I've come across a case in which the averaging includes padded values. Given a tensor X of some shape (batch_size, ..., features), there could be zero padded features to get the same shape.

How can I average the final dimension of X (the features) but only the non-zero entries? So, we divide by the sum by the number of non-zero entries.

Example input:

x = [[[[1,2,3], [2,3,4], [0,0,0]],
       [[1,2,3], [2,0,4], [3,4,5]],
       [[1,2,3], [0,0,0], [0,0,0]],
       [[1,2,3], [1,2,3], [0,0,0]]],
      [[[1,2,3], [0,1,0], [0,0,0]],
       [[1,2,3], [2,3,4], [0,0,0]],                                                         
       [[1,2,3], [0,0,0], [0,0,0]],                                                         
       [[1,2,3], [1,2,3], [1,2,3]]]]
# Desired output
y = [[[1.5 2.5 3.5]
      [2.  2.  4. ]
      [1.  2.  3. ]
      [1.  2.  3. ]]
     [[0.5 1.5 1.5]
      [1.5 2.5 3.5]
      [1.  2.  3. ]
      [1.  2.  3. ]]]
keras tensorflow python
2 Answers

A pure Keras solution counts the number of non-zero entries and then divides the sum accordingly. Here is a custom layer:

import keras.layers as L
import keras.backend as K

class NonZeroMean(L.Layer):
  """Compute mean of non-zero entries."""
  def call(self, x): 
    """Calculate non-zero mean."""
    # count the number of nonzero features, last axis
    nonzero = K.any(K.not_equal(x, 0.0), axis=-1)
    n = K.sum(K.cast(nonzero, 'float32'), axis=-1, keepdims=True)
    x_mean = K.sum(x, axis=-2) / n
    return x_mean

  def compute_output_shape(self, input_shape):
    """Collapse summation axis."""
    return input_shape[:-2] + (input_shape[-1],)

I suppose a condition needs to be added to check if all the features are zero and return zero, otherwise we get a division by zero error. Current example tested with:

# Dummy data
x = [[[[1,2,3], [2,3,4], [0,0,0]],
      [[1,2,3], [2,0,4], [3,4,5]],
      [[1,2,3], [0,0,0], [0,0,0]],
      [[1,2,3], [1,2,3], [0,0,0]]],
     [[[1,2,3], [0,1,0], [0,0,0]],
      [[1,2,3], [2,3,4], [0,0,0]],
      [[1,2,3], [0,0,0], [0,0,0]],
      [[1,2,3], [1,2,3], [1,2,3]]]]
x = np.array(x, dtype='float32')

# Example run
x_input = K.placeholder(shape=x.shape, name='x_input')
out = NonZeroMean()(x_input)
s = K.get_session()
print("INPUT:", x)
print("OUTPUT:", s.run(out, feed_dict={x_input: x}))
user_97589220

I accidentally encountered this question and noticed that the only solution is a bit overcomplicated. I will post an alternative solution for people who can potentially stumble here. In most cases, there would be no need to create a dedicated keras layer. One can do the following in TF 2.X:

# data prep
import tensorflow as tf
x = [[[[1, 2, 3], [2, 3, 4], [0, 0, 0]],
      [[1, 2, 3], [2, 0, 4], [3, 4, 5]],
      [[1, 2, 3], [0, 0, 0], [0, 0, 0]],
      [[1, 2, 3], [1, 2, 3], [0, 0, 0]]],
     [[[1, 2, 3], [0, 1, 0], [0, 0, 0]],
      [[1, 2, 3], [2, 3, 4], [0, 0, 0]],
      [[1, 2, 3], [0, 0, 0], [0, 0, 0]],
      [[1, 2, 3], [1, 2, 3], [1, 2, 3]]]]
x = tf.convert_to_tensor(x, dtype=tf.float32)

# solution
non_zero = tf.cast(x != 0, tf.float32)
y = tf.reduce_sum(x, axis=-2) / tf.reduce_sum(non_zero, axis=-2)

The operations performed are pretty much the same as in the present solution, only that it uses built-in TF functions. The array is first summed along the desired axis (second to last) and divided by the count of non-zero elements along it.

For this solution, if the count along the dimension is zero, the result will be NaN for that element.

user_65886850
Related