Skip to content

Instantly share code, notes, and snippets.

Last active March 30, 2021 16:50
  • Star 7 You must be signed in to star a gist
  • Fork 4 You must be signed in to fork a gist
Star You must be signed in to star a gist
Save raingo/a5808fe356b8da031837 to your computer and use it in GitHub Desktop.
multi dimensional softmax with tensorflow
import tensorflow as tf
Multi dimensional softmax,
refer to
compute softmax along the dimension of target
the native softmax only supports batch_size x dimension
def softmax(target, axis, name=None):
with tf.name_scope(name, 'softmax', values=[target]):
max_axis = tf.reduce_max(target, axis, keep_dims=True)
target_exp = tf.exp(target-max_axis)
normalize = tf.reduce_sum(target_exp, axis, keep_dims=True)
softmax = target_exp / normalize
return softmax
Copy link

Why the max_axis = tf.reduce_max(target, axis, keep_dims=True) ?

Copy link

N-McA commented Dec 7, 2017

For numerical stability; softmax(x+c) = softmax(x), but note that e^x gets big fast...

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment