Search code examples
pythonmachine-learningkerasneural-networksoftmax

Keras custom softmax layer: Is it possible to have output neurons set to 0 in the output of a softmax layer based on zeros as data in an input layer?

I have a neural network with 10 output neurons in the last layer using softmax activation. I also know exactly that based on the input values, certain neurons in the output layer shall have 0 values. So I have a special input layer of 10 neurons, each of them being either 0 or 1.

Would it be somehow possible to force let's say the output neuron no. 3 to have value = 0 if the input neuron no 3 is also 0?

action_input = Input(shape=(10,), name='action_input')
...

x = Dense(10,  kernel_initializer = RandomNormal(),bias_initializer = RandomNormal() )(x)
x = Activation('softmax')(x)

I know that there is a method via which I can mask out the results of the output layer OUTSIDE the neural network, and have all non zero related outputs reshaped (in order to have a total sum of 1). But I would like to solve this issue within the network and use it during the training of the network, too. Shall I use a custom layer for this?

Solution

  • At the end I came up with this code:

    from keras import backend as K
import tensorflow as tf
def mask_output2(x):
    inp, soft_out = x
    # add a very small value in order to avoid having 0 everywhere
    c = K.constant(0.0000001, dtype='float32', shape=(32, 13))
    y = soft_out + c

    y = Lambda(lambda x: K.switch(K.equal(x[0],0), x[1], K.zeros_like(x[1])))([inp, soft_out])
    y_sum =  K.sum(y, axis=-1)

    y_sum_corrected = Lambda(lambda x: K.switch(K.equal(x[0],0), K.ones_like(x[0]), x[0] ))([y_sum])

    y_sum_corrected = tf.divide(1,y_sum_corrected)

    y = tf.einsum('ij,i->ij', y, y_sum_corrected)
    return y