Keras: How to make custom layer with custom gradients?

[zgbkdlm]

Hi,

Are there any ways to manually control the update of the weights of customized layer in Keras?

for example, if I have a custom layer:

class MyLayer(Layer):

    def __init__(self, output_dim, **kwargs):
        self.output_dim = output_dim
        super(MyLayer, self).__init__(**kwargs)

    def build(self, input_shape):
        # Create a trainable weight variable for this layer.
        self.kernel = self.add_weight(name='kernel', 
                                      shape=(input_shape[1], self.output_dim),
                                      initializer='uniform',
                                      trainable=True)
        super(MyLayer, self).build(input_shape)  # Be sure to call this at the end

    def call(self, x):
        return K.dot(x, self.kernel)

    def compute_output_shape(self, input_shape):
        return (input_shape[0], self.output_dim)

My question is:

1. How exactly are the gradient calculated? If we have a custom layer, the derivation of backpropogation on this layer will be different from usual MLP or CNN layer. How exactly gradients of weights are calculated in custom layer? Are there any ways to customize the calculation of gradients of weights in custom layer?

2. If I understand correctly, for example, Tensorflow backend, it will automatically calculate the gradient according to the computation graph that defined. But for some iterative algorithm in the layer (filtering etc. ), I doubt it can automatically derive the gradient.

[volotat]

Same question. I have a model with two branches in it that merge at the end. I need to proceed my custom merge function in a way that it constantly return 0 and 1 as gradient for each branch respectively and I can't find any way how to do that.

[zgbkdlm]

@Mylittlerapture I think if you use tenforflow backend, check this tf.RegisterGradient. But I don't know if we could directly use this method in Keras.

[rezamm]

I have the same problem. Did you find any solutions on this recently?

[lodiam]

Same issue here. Has anyone come up with anything?

[volotat]

Same issue here. Has anyone come up with anything?

Yeah, I found solution eventually. But turns out it was not what I wanted to do, and I'm pretty sure same true for you. Anyway here is the solution:

def gradient_control(x, a):
    eps = 1-e4 #set epsilon between 0.1 and 1-e8
    return K.stop_gradient(x)* (1 - a * (1-eps)) +  x * (eps + a * (1 - eps))

'x' here is a tensor that you can get through Lambda layer and 'a' float between 0 and 1

[chengchengpei]

@Mylittlerapture I am curious where to add the function gradient_control()?

[volotat]

@chengchengpei inside your model in a layer after which you want to control gradient.
like that

layer = Dense(256, activation = 'relu')(layer)
layer = Lambda(lambda x: gradient_control(x, 0.5))(layer )

[THHHomas]

I need to define my own gradient as well, Is there any one can provide a method in the scope of keras? Thanks a lot.

[macarena1807]

Hi there, same problem here. I found this solution in StackOverflow (here), haven't tested yet.

I have implemented a custom layer without trainable parameters, but found out that a lot of TensorFlow and tensorflow.keras.backend functions are not differentiable (not defined gradients). If someone tries the previous answer of Djib2011 and report if it works that would be great. I'll try it next days.