Access output of intermediate layers in Tensor-flow 2.0 in eager mode

Question

I have CNN that I have built using on Tensor-flow 2.0. I need to access outputs of the intermediate layers. I was going over other stackoverflow questions that were similar but all had solutions involving Keras sequential model.

I have tried using model.layers[index].output but I get

Layer conv2d has no inbound nodes.

I can post my code here (which is super long) but I am sure even without that someone can point to me how it can be done using just Tensorflow 2.0 in eager mode.

Answer 1

I stumbled onto this question while looking for an answer and it took me some time to figure out as I use the model subclassing API in TF 2.0 by default (as in here https://www.tensorflow.org/tutorials/quickstart/advanced ). If somebody is in a similar situation, all you need to do is assign the intermediate output you want, as an attribute of the class. Then keep the test_step without the @tf.function decorator and create its decorated copy, say val_step, for efficient internal computation of validation performance during training. As a short example, I have modified a few functions of the tutorial from the link accordingly. I'm assuming we need to access the output after flattening.

def call(self, x):
    x = self.conv1(x)
    x = self.flatten(x)
    self.intermediate=x #assign it as an object attribute for accessing later
    x = self.d1(x)
return self.d2(x)

#Remove @tf.function decorator from test_step for prediction
def test_step(images, labels):
    predictions = model(images, training=False)
    t_loss = loss_object(labels, predictions)
    test_loss(t_loss)
    test_accuracy(labels, predictions)
    return

#Create a decorated val_step for object's internal use during training
@tf.function
def val_step(images, labels):
    return test_step(images, labels)

Now when you run model.predict() after training, using the un-decorated test step, you can access the intermediate output using model.intermediate which would be an EagerTensor whose value is obtained simply by model.intermediate.numpy(). However, if you don't remove the @tf_function decorator from test_step, this would return a Tensor whose value is not so straightforward to obtain.

Answer 2

Thanks for answering my earlier question. I wrote this simple example to illustrate how what you're trying to do might be done in TensorFlow 2.x, using the MNIST dataset as the example problem.

The gist of the approach:

1. Build an auxiliary model ( aux_model in the example below), which is so-called "functional model" with multiple outputs. The first output is the output of the original model and will be used for loss calculation and backprop, while the remaining output(s) are the intermediate-layer outputs that you want to access.

2. Use tf.GradientTape() to write a custom training loop and expose the detailed gradient values on each individual variable of the model. Then you can pick out the gradients that are of interest to you. This requires that you know the ordering of the model's variables. But that should be relatively easy for a sequential model.

import tensorflow as tf

(x_train, y_train), (_, _) = tf.keras.datasets.mnist.load_data()

# This is the original model.
model = tf.keras.Sequential([
    tf.keras.layers.Flatten(input_shape=[28, 28, 1]),
    tf.keras.layers.Dense(100, activation="relu"),
    tf.keras.layers.Dense(10, activation="softmax")])

# Make an auxiliary model that exposes the output from the intermediate layer
# of interest, which is the first Dense layer in this case.
aux_model = tf.keras.Model(inputs=model.inputs,
                           outputs=model.outputs + [model.layers[1].output])

# Define a custom training loop using `tf.GradientTape()`, to make it easier
# to access gradients on specific variables (the kernel and bias of the first
# Dense layer in this case).
cce = tf.keras.losses.CategoricalCrossentropy()
optimizer = tf.optimizers.Adam()
with tf.GradientTape() as tape:
  # Do a forward pass on the model, retrieving the intermediate layer's output.
  y_pred, intermediate_output = aux_model(x_train)
  print(intermediate_output)  # Now you can access the intermediate layer's output.

  # Compute loss, to enable backprop.
  loss = cce(tf.one_hot(y_train, 10), y_pred)

# Do backprop. `gradients` here are for all variables of the model.
# But we know we want the gradients on the kernel and bias of the first
# Dense layer, which happens to be the first two variables of the model.
gradients = tape.gradient(loss, aux_model.variables)

# This is the gradient on the first Dense layer's kernel.
intermediate_layer_kerenl_gradients = gradients[0]
print(intermediate_layer_kerenl_gradients)

# This is the gradient on the first Dense layer's bias.
intermediate_layer_bias_gradients = gradients[1]
print(intermediate_layer_bias_gradients)

# Update the variables of the model.
optimizer.apply_gradients(zip(gradients, aux_model.variables))

Answer 3

The most straightforward solution would go like this:

mid_layer = model.get_layer("layer_name")

you can now treat the "mid_layer" as a model, and for instance:

mid_layer.predict(X)

Oh, also, to get the name of a hidden layer, you can use this:

model.summary()

this will give you some insights about the layer input/output as well.