Tensorflow median value

For calculating median of an array with tensorflow you can use the percentile function, since the 50th percentile is the median.

import tensorflow as tf
import tensorflow_probability as tfp
import numpy as np 

np.random.seed(0)   
x = np.random.normal(3.0, .1, 100)

median = tfp.stats.percentile(x, 50.0, interpolation='midpoint')

tf.Session().run(median)

The code above is equivalent to np.percentile(x, 50, interpolation='midpoint').

We can modify BlueSun's solution to be much faster on GPUs:

def get_median(v):
    v = tf.reshape(v, [-1])
    m = v.get_shape()[0]//2
    return tf.reduce_min(tf.nn.top_k(v, m, sorted=False).values)

This is as fast as (in my experience) using tf.contrib.distributions.percentile(v, 50.0), and returns one of the actual elements.

edit: This answer is outdated, use Lucas Venezian Povoa's solution instead. It is simpler and faster.

You can calculate the median inside tensorflow using:

def get_median(v):
    v = tf.reshape(v, [-1])
    mid = v.get_shape()[0]//2 + 1
    return tf.nn.top_k(v, mid).values[-1]

If X is already a vector you can skip the reshaping.

If you care about the median value being the mean of the two middle elements for vectors of even size, you should use this instead:

def get_real_median(v):
    v = tf.reshape(v, [-1])
    l = v.get_shape()[0]
    mid = l//2 + 1
    val = tf.nn.top_k(v, mid).values
    if l % 2 == 1:
        return val[-1]
    else:
        return 0.5 * (val[-1] + val[-2])