Tensorflow，在RNN中保存状态的最佳方法？

Question

我现在有下面的代码用于tensorflow中一系列链接在一起的rnn。我没有使用MultiRNN，因为我稍后要对每个层的输出做一些事情。

 for r in range(RNNS):
    with tf.variable_scope('recurent_%d' % r) as scope:
        state = [tf.zeros((BATCH_SIZE, sz)) for sz in rnn_func.state_size]
        time_outputs = [None] * TIME_STEPS

        for t in range(TIME_STEPS):
            rnn_input = getTimeStep(rnn_outputs[r - 1], t)
            time_outputs[t], state = rnn_func(rnn_input, state)
            time_outputs[t] = tf.reshape(time_outputs[t], (-1, 1, RNN_SIZE))
            scope.reuse_variables()
        rnn_outputs[r] = tf.concat(1, time_outputs)

目前我有固定的时间步数。不过，我希望将其更改为只有一个timestep，但记住批处理之间的状态。因此，我需要为每个层创建一个状态变量，并为其分配每个层的最终状态。像这样的东西。

for r in range(RNNS):
    with tf.variable_scope('recurent_%d' % r) as scope:
        saved_state = tf.get_variable('saved_state', ...)
        rnn_outputs[r], state = rnn_func(rnn_outputs[r - 1], saved_state)
        saved_state = tf.assign(saved_state, state)

然后，对于每个层，我需要评估sess.run函数中保存的状态，并调用我的训练函数。我需要对每个rnn层都这样做。这看起来有点麻烦。我需要跟踪每个保存的状态并在运行时对其进行评估。另外，run需要将状态从我的GPU复制到主机内存，这将是低效和不必要的。有更好的办法吗？

Answer 1

我现在使用tf.control_依赖项保存RNN状态。这是一个例子。

 saved_states = [tf.get_variable('saved_state_%d' % i, shape = (BATCH_SIZE, sz), trainable = False, initializer = tf.constant_initializer()) for i, sz in enumerate(rnn.state_size)]
        W = tf.get_variable('W', shape = (2 * RNN_SIZE, RNN_SIZE), initializer = tf.truncated_normal_initializer(0.0, 1 / np.sqrt(2 * RNN_SIZE)))
        b = tf.get_variable('b', shape = (RNN_SIZE,), initializer = tf.constant_initializer())

        rnn_output, states = rnn(last_output, saved_states)
        with tf.control_dependencies([tf.assign(a, b) for a, b in zip(saved_states, states)]):
            dense_input = tf.concat(1, (last_output, rnn_output))

        dense_output = tf.tanh(tf.matmul(dense_input, W) + b)
        last_output = dense_output + last_output

我只是确保我的图形的一部分依赖于保存状态。