Custom Lambda layer Keras outputs predictions. I get 'An operation has `None` for gradient' error

Question

I have a Lambda layer that takes input from previous layer, makes some preprocessing. Output of the Lambda layer is a prediction, and keras.losses.mean_squared_error is used.

#model
inputs = Input(shape=(len_train_data_columns,))
dense1 = Dense(777, activation='relu')(inputs)
dropout1 = Dropout(0.4)(dense1)
softmax = Dense(3, activation='softmax')(dropout1)
predictions = Lambda(custom_layer,
            output_shape= custom_layer_output_shape,
            arguments= {'experiment_config': experiment_config},
               )(softmax)

model = Model(inputs = inputs, outputs = predictions)
model.compile(optimizer=adam(lr=0.0001), loss=keras.losses.mean_squared_error, metrics=keras.losses.mean_squared_error)

Output of Lambda is (num_rows, 1).

I get ValueError: An operation has None for gradient. Please make sure that all of your ops have a gradient defined (i.e. are differentiable). Common ops without gradient: K.argmax, K.round, K.eval.

Why is that? I use Keras loss...

Answer 1

I think it comes from your 'custom_layer'. Not all functions defined in keras backend are differentiable (as the error says, for example K.argmax) if you use these functions on your layer and they don't have a defined gradient, it will rise this error. You have to make sure that you use only functions that has a gradient associated. Although there is no an official list of these functions, basic operations (+/-*) are, for example equal_toor greaterare not differentiable.

Answer 2

You could turn off back propagation for some nodes, layers or functions, but I would not recommend that. You can also define a custom gradient for your whole custom layer, if you know what it should be. For example if it's an implementation of a well-known function you can just derive manually the gradient and then implement it as a custom gradient.

If not, you should fight a little bit and try to re-implement the custom layer using TensorFlow differentiable ops.