Suppose you have a Keras NN model, how can you stop the gradient in the backpropagation after a certain layer?
I.e., if we have a model with two outputs:
input_layer = Input(shape=(10,10,3))
x = Convolution2D(...)(input_layer)
x = Activation('relu')(x)
x = Flatten()(x)
x_1 = Dense(64)(x)
x_1 = Dense(32)(x_1)
x_1 = Dense(2)(x_1)
x_2 = Dense(64)(x)
x_2 = Dense(32)(x_2)
x_2 = Dense(2)(x_2)
model = Model(inputs=input_layer, outputs=[x_1, x_2])
How can the gradient of output x_1 be stopped after the x_1 = Dense(64)(x) layer, such that it doesn't get counted in for the weight update in the convolutional layer?
Based on the answer in Stopping Gradient back prop through a particular layer in keras, I'd have added a lambda layer before the x_1 dense layer, but I'm not really sure:
x_1 = Dense(64)(x)
x_1_stop_grad = Lambda(lambda x: K.stop_gradient(x))(x_1)
x_1 = Dense(32)(x_1)
x_1 = Dense(2)(x_1)
Do I have to add the lambda layer before or after the first dense x_1 layer?
Since the gradient is flowing backwards through the network, you need to add the gradient stop layer directly after the layer, where no gradient should arrive.
I.e.
from keras import ops
# weights in x should not be updated by gradients from x_1
x = Convolution2D(...)(input_layer)
x_1_stop_grad = Lambda(lambda x: ops.stop_gradient(x))(x)
x_1 = Dense(64)(x_1_stop_grad)
x_1 = Dense(32)(x_1)