Skip to content

Instantly share code, notes, and snippets.

@raingo
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 https://github.com/tensorflow/tensorflow/issues/210
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
@cancan101
Copy link

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

@N-McA
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