How to specify the axis when using the softmax activation in a Keras layer?

Question

The Keras docs for the softmax Activation states that I can specify which axis the activation is applied to. My model is supposed to output an n by k matrix M where Mij is the probability that the ith letter is symbol j.

n = 7 # number of symbols in the ouput string (fixed)
k = len("0123456789") # the number of possible symbols

model = models.Sequential()
model.add(layers.Dense(16, activation='relu', input_shape=((N,))))
...
model.add(layers.Dense(n * k, activation=None))
model.add(layers.Reshape((n, k)))

model.add(layers.Dense(output_dim=n, activation='softmax(x, axis=1)'))

The last line of code doesn't compile as I don't know how to correctly specify the axis (the axis for k in my case) for the softmax activation.

Answer 1

You must use an actual function there, not a string.

Keras allows you to use a few strings for convenience.

The activation functions can be found in keras.activations, and they're listed in the help file.

from keras.activations import softmax

def softMaxAxis1(x):
    return softmax(x,axis=1)

..... 
......
model.add(layers.Dense(output_dim=n, activation=softMaxAxis1))

---

Or even a custom axis:

def softMaxAxis(axis):
    def soft(x):
        return softmax(x,axis=axis)
    return soft

...
model.add(layers.Dense(output_dim=n, activation=softMaxAxis(1)))