用于逐元素乘法的 Lambda 层在第一次模型更新后给出 NaN。 (凯拉斯)

Question

My model is written in Keras. It has multiple inputs, one of which is to be multiplied with the outputs of the penultimate Dense layer before the logits are fed into the softmax. This element wise multiplication is carried out by means of a Lambda Layer.

def mul(x, mask):
    output = x*mask
    return output

logits = Lambda(lambda x: mul(x, input_2))(dense_output) # gives nan after first update

After the model is updated for the first time, the Lambda Layer starts giving NaN as output .
This does not happen if I have a constant vector defined within the computational graph being multiplied with the output of the dense layer.

logits = dense_output * [1, 1, 1, 1, 1, -100, 1, -100, 1, 1] # does not give nan

I have tried using the Multiply Layer provided by Keras as well. But this too throws NaN after the first update. Here is a snippet for the same:

logits = Multiply()([dense_output, input_2]) # gives nan after first update

I basically want to mask certain output states by the means of this multiplication with the input, but can't do it if the layer keeps giving NaN as output.
Is there any way to solve this? Any and all help will be appreciated!

Answer 1

I am not sure how this works, but I have found a solution to this. It seems that adding a Dense layer in front of the Multiply layer solves the problem. It doesn't matter if the Dense layer is trainable or not. Here is the code:

logits = Multiply()([dense_output, input_2])
initializer = tf.keras.initializers.Identity()
masked_actions = Dense(num_actions, use_bias=False, 
                       trainable=False)(logits) #returns the same logits

The model updates now work as expected without throwing any NaN ValueErrors.