Extracting hidden state vectors of an LSTM model at each time step

user_1660449 0
Viewed 2016

I'm using keras with the tensorflow backend.

I have a trained LSTM model whose hidden state vector I would like to extract at each timestep.

What is the best way to do this in keras?

keras tensorflow
1 Answers

The function handling whether all hidden state vectors are returned is Recurrent.call() (it has been renamed to RNN.call() in the latest version). It checks the parameter return_sequences to make the decision.

When the backend function K.rnn() is called in this function:

last_output, outputs, states = K.rnn(self.step,
                                     preprocessed_input,
                                     initial_state,
                                     go_backwards=self.go_backwards,
                                     mask=mask,
                                     constants=constants,
                                     unroll=self.unroll,
                                     input_length=input_shape[1])

...

if self.return_sequences:
    output = outputs
else:
    output = last_output

The tensor outputs is what you want. You can get this tensor by calling Recurrent.call() again, but with return_sequences=True. This should do no harm to your trained LSTM model (at least in current Keras).

Here's a toy Bi-LSTM model demonstrating this method:

input_tensor = Input(shape=(None,), dtype='int32')
embedding = Embedding(10, 100, mask_zero=True)(input_tensor)
hidden = Bidirectional(LSTM(10, return_sequences=True))(embedding)
hidden = Bidirectional(LSTM(10, return_sequences=True))(hidden)
hidden = Bidirectional(LSTM(2))(hidden)
out = Dense(1, activation='sigmoid')(hidden)
model = Model(input_tensor, out)

First, set return_sequences to True for the last LSTM layer (since you're using a Bidirectional wrapper, you have to set forward_layer and backward_layer also):

target_layer = model.layers[-2]
target_layer.return_sequences = True
target_layer.forward_layer.return_sequences = True
target_layer.backward_layer.return_sequences = True

Now by calling this layer again, the tensor containing hidden vectors at all time steps will be returned (there will be a side-effect of creating an additional inbound node, but it shouldn't affect the prediction).

outputs = target_layer(target_layer.input)
m = Model(model.input, outputs)

You can get the hidden vectors by, e.g., calling m.predict(X_test).

X_test = np.array([[1, 3, 2, 0, 0]])
print(m.predict(X_test))

[[[ 0.00113332 -0.0006666   0.00428438 -0.00125567]
  [ 0.00106074 -0.00041183  0.00383953 -0.00027285]
  [ 0.00080892  0.00027685  0.00238486  0.00036328]
  [ 0.00080892  0.00027685  0.          0.        ]
  [ 0.00080892  0.00027685  0.          0.        ]]]

As you can see, hidden vectors of all 5 time steps are returned, and the last 2 time steps are properly masked.

user_15314630
Related