基于TensorFlow后端的Keras GRU模型预测时隐藏状态的保存与恢复

I built a stateful GRU model using Keras with TensorFlow backend as follows:

X_input = Input(name="x_input", batch_shape=(BATCH_SIZE, 1, len(alphabet_dict)))
GRU_layer, hs = GRU(32, return_sequences=True, dropout=0.25, recurrent_dropout=0.25, stateful=True, return_state=True)(X_input)
GRU_layer2, hs2 = GRU(32, return_sequences=True, dropout=0.25, recurrent_dropout=0.25, stateful=True, return_state=True)(GRU_layer)
y = Dense(len(alphabet_dict), activation="softmax", name="y")(GRU_layer2)
model = Model(inputs=[X_input], outputs=[y, hs, hs2])
model.compile(loss={"y": categorical_crossentropy}, optimizer="adam", metrics=["acc"])

After training for several epochs, I saved the model. Then in the prediction script, I restored the model from the checkpoint and rebuilt an independent model with batch_size=1:

with tf.device('/cpu:0'):
    model = load_model(filepath=model_file, compile=True)
    old_weights = model.get_weights()
    del model

X_input = Input(name="x_input", batch_shape=(1, 1, len(alphabet_dict)))
GRU_layer, hs = GRU(32, return_sequences=True, dropout=0.25, recurrent_dropout=0.25, stateful=True, return_state=True)(X_input)
GRU_layer2, hs2 = GRU(32, return_sequences=True, dropout=0.25, recurrent_dropout=0.25, stateful=True, return_state=True)(GRU_layer)
y = Dense(len(alphabet_dict), activation="softmax", name="y")(GRU_layer2)
model = Model(inputs=[X_input], outputs=[y, hs, hs2])
model.compile(loss={"y": categorical_crossentropy}, optimizer="adam", metrics=["acc"])
model.set_weights(old_weights)

How can I implement saving and restoring the hidden states of the GRU layers in this prediction workflow?

Great question—working with stateful RNNs means dealing with runtime hidden states that aren’t saved alongside model weights, so you need a separate process to handle them. Here’s how to adapt your workflow to save and restore these states properly:

1. Save Hidden States After Training #

When your stateful model finishes training, extract the internal hidden states from each GRU layer and save them to a file (we’ll use NumPy for simplicity):

# After training completes, fetch the GRU layers from your trained model
gru_layer1 = model.get_layer(index=1)  # Adjust index if your layer order differs
gru_layer2 = model.get_layer(index=2)

# Extract the hidden states: get_state() returns a list (GRU has one state tensor)
state1 = gru_layer1.get_state()[0]
state2 = gru_layer2.get_state()[0]

# Save the states to a compressed NPZ file
import numpy as np
np.savez('gru_hidden_states.npz', state1=state1, state2=state2)

Pro tip: Instead of relying on layer indices, you can name your GRU layers during creation (e.g., GRU(32, name="gru_1", ...)) and use model.get_layer("gru_1") for more reliable access.

2. Restore Hidden States in Your Prediction Model #

Once you’ve built your batch_size=1 model and loaded the trained weights, load the saved states and adjust them to match the new batch size, then set them on the corresponding layers:

# Load the saved hidden states
loaded_states = np.load('gru_hidden_states.npz')
saved_state1 = loaded_states['state1']
saved_state2 = loaded_states['state2']

# Fetch the GRU layers from your new prediction model
new_gru1 = model.get_layer(index=1)
new_gru2 = model.get_layer(index=2)

# Adjust the batch dimension: original state shape is (BATCH_SIZE, 32)
# We need to slice it to (1, 32) for our prediction batch size
# Pick the first sample's state (or any sample relevant to your use case)
adjusted_state1 = saved_state1[0:1, :]
adjusted_state2 = saved_state2[0:1, :]

# Set the state on the layers (pass as a list, since get_state() returns a list)
new_gru1.set_state([adjusted_state1])
new_gru2.set_state([adjusted_state2])

3. Maintain State During Sequential Predictions (Optional) #

If you’re making a sequence of predictions (e.g., generating text character by character), you’ll want to save the updated state after each prediction to keep continuity:

# Run a prediction step
prediction, updated_state1, updated_state2 = model.predict(your_input_data)

# Save the fresh states for the next prediction step
np.savez('current_gru_states.npz', state1=updated_state1, state2=updated_state2)

Before your next prediction, load this file and repeat the state-setting process from step 2.

Key Notes to Avoid Issues #

• State Shape Matching: Double-check the shape of your saved states with print(saved_state1.shape)—for GRU, it should be (batch_size, units) (in your case, (BATCH_SIZE, 32)). When adjusting to batch_size=1, slice along the first dimension.
• Layer Configuration Consistency: Ensure your prediction model’s GRU layers have identical configurations (units, stateful=True, return_state=True, etc.) to the training model. Mismatches will cause errors when setting states.
• Batch Element Relevance: Stateful RNNs store a separate state for each element in the batch. When switching to batch_size=1, pick the batch element’s state that’s relevant to your prediction task (e.g., the first element if you trained on ordered sequences).

内容的提问来源于stack exchange，提问作者Bocharick