How to get outputs of all intermediate layers of a deep RNN in Keras

Question

I am trying to obtain the outputs of every layer of a custom RNN model in Keras. The code for the model is given below.

from tensorflow.python import keras
from tensorflow.python.keras import backend as K
from tensorflow.python.keras.layers import RNN, Dense, Activation
from tensorflow.python.keras.models import Model
import numpy as np


class MinimalRNNCell(keras.layers.Layer):

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

    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]


cells = [MinimalRNNCell(32), MinimalRNNCell(64)]
x = keras.Input((None, 257))
y = RNN(cells)(x)
out = Dense(257, name='fin_dense')(y)
out = Activation('sigmoid', name='out_layer')(out)

model = Model(inputs=x, outputs=out)

I can get the predictions using,

in_test = np.random.randn(1, 3, 257)
mod_out = model.predict(in_test)

But I would like the outputs of every layer and when I try using the getLayerOutputs function which works for all other Keras models

def getLayerOutputs(model, input_data, learning_phase=1):
    outputs = [layer.output for layer in model.layers[1:]] # exclude Input
    layers_fn = K.function([model.input, K.learning_phase()], outputs)
    return layers_fn([input_data, learning_phase])

layer_outs = getLayerOutputs(model, in_test)

I get the output of the entire RNN layer and not each cell within it, how do I get the output of each cell with the RNN?

Answer 1

This should do it on tensorflow 1.x (test version 1.12.0):

sess = K.get_session()

outs = dict() 
for i in range(1, len(model.layers)): 
    outs[model.layers[i].name] = sess.run([model.layers[i].output], feed_dict={model.input:in_test})  

(outs has the intermediate results for each layer)

If you are using TensorFlow 2 (test version 2.0.0-rc1):

outs = dict() 
for i in range(1, len(model.layers)): 
    outs[model.layers[i].name] = Model(x, model.layers[i].output).predict(in_test)

Answer 2

By changing return_state flag as True like so

y = RNN(cells, return_state=True)(x)

We get output states of every cell of the RNN as a list