Search code examples
tensorflowrecurrent-neural-networkrestore

Tensorflow 1.0 : Retreive hidden state from restored RNN


I'd like to restore an RNN and get the hidden state.

I do something like that to save the RNN:

loc="path/to/save/rnn"
with tf.variable_scope("lstm") as scope:
    outputs, state = tf.nn.dynamic_rnn(..)
sess = tf.Session()
sess.run(tf.global_variables_initializer())
saver = tf.train.Saver()
save_path = saver.save(sess,loc)

Now I want to retreive state.

graph = tf.Graph()
sess = tf.Session(graph=graph)
with graph.as_default():
      saver = tf.train.import_meta_graph(loc + '.meta', clear_devices=True)
      saver.restore(sess, loc)
      state= ...

Solution

  • You can add the state tensor to a graph collection, which is basically a key value store to track tensors, using tf.add_to_collection and retrieve it later using tf.get_collection. For example:

    loc="path/to/save/rnn"
    with tf.variable_scope("lstm") as scope:
        outputs, state = tf.nn.dynamic_rnn(..)
        tf.add_to_collection('state', state)
    
    
    graph = tf.Graph()
    with graph.as_default():
          saver = tf.train.import_meta_graph(loc + '.meta', clear_devices=True)
          state = tf.get_collection('state')[0]  # Note: tf.get_collection returns a list.