How to get loss gradient wrt internal layer output in tensorflow 2?

Question

I would like to get gradient of the model's loss function with respect to specific layer's output during training. What I want to do with it next, is using a value of that gradient to modify something in layer in the next learning epoch. So how to obtain that gradient?

Here's a minimal example. MinimalRNNCell code is copied from TensorFlow's website and toy data is provided only to reproduce the behavior.

import tensorflow as tf 
from tensorflow.keras.layers import RNN, SimpleRNNCell, SimpleRNN, Layer, Dense, AbstractRNNCell
from tensorflow.keras import Model
import numpy as np
import tensorflow.keras.backend as K


class MinimalRNNCell(AbstractRNNCell):

    def __init__(self, units, **kwargs):
      self.units = units
      super(MinimalRNNCell, self).__init__(**kwargs)

    @property
    def state_size(self):
      return self.units

    def build(self, input_shape):
      self.kernel = self.add_weight(shape=(input_shape[-1], self.units),
                                    initializer='uniform',
                                    name='kernel')
      self.recurrent_kernel = self.add_weight(
          shape=(self.units, self.units),
          initializer='uniform',
          name='recurrent_kernel')
      self.built = True

    def call(self, inputs, states):
      prev_output = states[0]
      h = K.dot(inputs, self.kernel)
      output = h + K.dot(prev_output, self.recurrent_kernel)
      return output, output


class MyModel(Model):
    def __init__(self, size):
        super(MyModel, self).__init__()
        self.minimalrnn=RNN(MinimalRNNCell(size), name='minimalrnn')
        self.out=Dense(4)

    def call(self, inputs):
        out=self.minimalrnn(inputs)
        out=self.out(out)
        return out


x=np.array([[[3.],[0.],[1.],[2.],[3.]],[[3.],[0.],[1.],[2.],[3.]]])
y=np.array([[[0.],[1.],[2.],[3.]],[[0.],[1.],[2.],[3.]]])

model=MyModel(2)
model.compile(optimizer='sgd', loss='mse')
model.fit(x,y,epochs=10, batch_size=1, validation_split=0.2)

Now I want to get gradient of output of MyModel's minimalrnn layer (after every batch of data).

How to do this? I suppose I can try with GradientTape watching model.get_layer('minimalrnn').output, but I need more learning resources or examples.

EDIT

I used GradientTape as in code provided by Tiago Martins Peres, but I specifically want to obtain gradient wrt layer output, and I'm still not able to achieve that.

Now after class definitions my code looks like this:

x=np.array([[[3.],[0.],[1.],[2.],[3.]],[[3.],[0.],[1.],[2.],[3.]]])
y=np.array([[0., 1., 2., 3.],[0., 1., 2., 3.]])

model=MyModel(2)

#inputs = tf.keras.Input(shape=(2,5,1))
#model.call(x)

def gradients(model, inputs, targets):
    with tf.GradientTape() as tape:
        tape.watch(model.get_layer('minimalrnn').output)
        loss_value = loss_fn(model, inputs, targets)
    return tape.gradient(loss_value, model.trainable_variables)

def loss_fn(model, inputs, targets):
    error = model(inputs) - targets
    return tf.reduce_mean(tf.square(error))

optimizer = tf.keras.optimizers.Adam(learning_rate=0.01)
print("Initial loss: {:.3f}".format(loss_fn(model, x, y)))
for i in range(10):
    grads = gradients(model, x, y)
    optimizer.apply_gradients(zip(grads, model.trainable_variables))
    print("Loss at step {:03d}: {:.3f}".format(i, loss_fn(model, x, y)))
print("Final loss: {:.3f}".format(loss_fn(model, x, y)))

As you can see I added tape.watch in gradients function definition, because I want to watch layer output. However I'm getting error:

Traceback (most recent call last):
  File "/home/.../test2.py", line 73, in <module>
    grads = gradients(model, x, y)
  File "/home/.../test2.py", line 58, in gradients
    print(model.get_layer('minimalrnn').output)
  File "/home/.../.venv/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/base_layer.py", line 1553, in output
    raise AttributeError('Layer ' + self.name + ' has no inbound nodes.')
AttributeError: Layer minimalrnn has no inbound nodes.

I also tried to call model on Input with specified size (commented lines), according to answer to this: Accessing layer's input/output using Tensorflow 2.0 Model Sub-classing. It didn't help. Specifying input shape in model's init function, like below, also doesn't help - still the same error.

self.minimalrnn=RNN(MinimalRNNCell(size), name='minimalrnn', input_shape=(2,5,1))

Answer 1

Ok, so one answer that I finally found is hidden here: https://stackoverflow.com/a/56567364/4750170. I can even use subclassed model with this.

Additionally problem with AttributeError is strange, because when I used Sequential instead of subclassing Model, AttributeError magically disappeared, maybe it's connected with this issue https://github.com/tensorflow/tensorflow/issues/34834?

Still, I'd like to know why I can't just pass the layer's output as a second argument to tape.gradient.

Answer 2

Yes you can use GradientTape. The purpose of tf.GradientTape is to record operations for automatic differentiation or for computing the gradient of an operation or computation with respect to its input variables.

According to What's New in TensorFlow 2.0, to first implement the simple training of a model with tf.GradientTape, call the forward pass on the input tensor inside the tf.GradentTape context manager and then compute the loss function. This ensures that all of the computations will be recorded on the gradient tape.

Then, compute the gradients with regard to all of the trainable variables in the model. Once the gradients are computed, any desired gradient clipping, normalization, or transformation can be performed before passing them to the optimizer to apply them to the model variables. Take a look at the following example:

NUM_EXAMPLES = 2000

input_x = tf.random.normal([NUM_EXAMPLES])
noise = tf.random.normal([NUM_EXAMPLES])
input_y = input_x * 5 + 2 + noise

def loss_fn(model, inputs, targets):
  error = model(inputs) - targets
  return tf.reduce_mean(tf.square(error))

def gradients(model, inputs, targets):
  with tf.GradientTape() as tape:
    loss_value = loss_fn(model, inputs, targets)
  return tape.gradient(loss_value, model.trainable_variables)

model = tf.keras.Sequential(tf.keras.layers.Dense(1))
optimizer = tf.keras.optimizers.Adam(learning_rate=0.01)
print("Initial loss: {:.3f}".format(loss_fn(model, input_x, input_y)))
for i in range(500):
  grads = gradients(model, input_x, input_y)
  optimizer.apply_gradients(zip(grads, model.trainable_variables))
  if i % 20 == 0:
    print("Loss at step {:03d}: {:.3f}".format(i, loss_fn(model, input_x, input_y)))
print("Final loss: {:.3f}".format(loss(model, input_x, input_y)))
print("W = {}, B = {}".format(*model.trainable_variables))