I am currently trying to programmatically visualize the way deep learning approaches learn/train on my problems and e.g. stumbled upon the paper "Visualizing and Understanding Recurrent Networks" by Andrej Karpathy, Justin Johnson, and Li Fei-Fei (e.g. https://arxiv.org/abs/1506.02078).
In figure 2 they calculate the tanh of the memory/cell states (c) of their LSTMs at each timestep during prediction.
I now want to do the same with my LSTMs in Keras (or PyTorch, but I'm quite new to their API). The API allows me to get the hidden states (h) for each timestep - using the parameter return_sequences=True, but only the last cell state - using the parameter return_state=True.
With my understanding of the Keras API I currently create two models:
inputs = Input(shape=(MAX_LEN,))
embd = Embedding(input_dim=MAX_FEATURES, output_dim=128, input_length=MAX_LEN)(inputs)
lstm = LSTM(units=128, return_sequences=True)(embd)
# 'ignore' return_sequences=True in training model
flt = Flatten()(lstm)
drp = Dropout(0.5)(flt)
dns = Dense(N_CLASSES)(drp)
activ = Activation('softmax')(dns)
training_model = Model([inputs], [activ])
state_model = Model([inputs], [activ, lstm])
training_model.compile(loss='categorical_crossentropy', optimizer='rmsprop')
training_model.fit(X_train, y_train)
predictions, hidden_states = state_model.predict(X_test)
How do I do the same for the cell state at each timestep (character in my case) - preferably without having to override/reimplement 'major' parts of Keras'/PyTorch's source code? I found this answer, but haven't gotten it to work and as I don't have enough reputation I can't add any questions in the comments (I saw the other answer but I want to prevent a 'complete' reimplementation of Keras' LSTMCell).
EDIT #1: As I'm quite new to PyTorch's API - how would I go about there?