Sonnet 中的 LSTM 时间步长

Question

我目前正在努力学习Sonnet。

我的网络（不完整，问题基于此）：

class Model(snt.AbstractModule):

    def __init__(self, name="LSTMNetwork"):
        super(Model, self).__init__(name=name)
        with self._enter_variable_scope():
            self.l1 = snt.LSTM(100)
            self.l2 = snt.LSTM(100)
            self.out = snt.LSTM(10)

    def _build(self, inputs):

        # 'inputs' is of shape (batch_size, input_length)
        # I need it to be of shape (batch_size, sequence_length, input_length)

        l1_state = self.l1.initialize_state(np.shape(inputs)[0]) # init with batch_size
        l2_state = self.l2.initialize_state(np.shape(inputs)[0]) # init with batch_size
        out_state = self.out.initialize_state(np.shape(inputs)[0])

        l1_out, l1_state = self.l1(inputs, l1_state)
        l1_out = tf.tanh(l1_out)
        l2_out, l2_state = self.l2(l1_out, l2_state)
        l2_out = tf.tanh(l2_out)
        output, out_state = self.out(l2_out, out_state)
        output = tf.sigmoid(output)

        return output, out_state

在其他框架（例如 Keras）中，LSTM 输入的格式为 (batch_size, sequence_length, input_length)。

但是，Sonnet 文档指出 Sonnet 的 LSTM 的输入格式为 (batch_size, input_length)。

如何将它们用于顺序输入？

到目前为止，我已经尝试在 _build 中使用 for 循环，迭代每个时间步长，但这会产生看似随机的输出。

我在 Keras 中尝试过相同的架构，运行时没有任何问题。

我正在以 Eager 模式执行，使用 GradientTape 进行训练。

Answer 1

我们通常在 Sonnet 中编写 RNN 以在单个时间步的基础上工作，对于强化学习，您通常需要运行一个时间步来选择一个动作，如果没有该动作，您将无法获得下一个观察结果（以及下一个输入时间步长）来自环境。使用tf.nn.dynamic_rnn（见下文）在序列上展开单个时间步长模块很容易。我们还有一个包装器，它负责每个时间步长组成几个 RNN 核心，我相信这就是您想要做的。这样做的好处是DeepCore 对象支持dynamic_rnn 所需的启动状态方法，因此它的API 与LSTM 或任何其他单时间步模块兼容。

你想做的事情应该可以这样实现：

# Create a single-timestep RNN module by composing recurrent modules and
# non-recurrent ops.
model = snt.DeepRNN([
    snt.LSTM(100),
    tf.tanh,
    snt.LSTM(100),
    tf.tanh,
    snt.LSTM(100),
    tf.sigmoid
], skip_connections=False)

batch_size = 2
sequence_length = 3
input_size = 4

single_timestep_input = tf.random_uniform([batch_size, input_size])
sequence_input = tf.random_uniform([batch_size, sequence_length, input_size])

# Run the module on a single timestep
single_timestep_output, next_state = model(
    single_timestep_input, model.initial_state(batch_size=batch_size))

# Unroll the module on a full sequence
sequence_output, final_state = tf.nn.dynamic_rnn(
    core, sequence_input, dtype=tf.float32)

有几点需要注意 - 如果您还没有，请查看存储库中的 RNN example，因为这显示了围绕一个非常相似的模型设置的完整图形模式训练过程。

其次，如果您最终需要实现DeepRNN 允许的更复杂的模块，那么将循环状态线程化进出模块非常重要。在您的示例中，您在内部创建输入状态，并且 l1_state 和 l2_state 作为输出被有效地丢弃，因此无法正确训练。如果 DeepRNN 不可用，您的模型将如下所示：

class LSTMNetwork(snt.RNNCore):  # Note we inherit from the RNN-specific subclass
  def __init__(self, name="LSTMNetwork"):
    super(Model, self).__init__(name=name)
    with self._enter_variable_scope():
      self.l1 = snt.LSTM(100)
      self.l2 = snt.LSTM(100)
      self.out = snt.LSTM(10)

  def initial_state(self, batch_size):
    return (self.l1.initial_state(batch_size),
            self.l2.initial_state(batch_size),
            self.out.initial_state(batch_size))

  def _build(self, inputs, prev_state):

    # separate the components of prev_state
    l1_prev_state, l2_prev_state, out_prev_state = prev_state

    l1_out, l1_next_state = self.l1(inputs, l1_prev_state)
    l1_out = tf.tanh(l1_out)
    l2_out, l2_next_state = self.l2(l1_out, l2_prev_state)
    l2_out = tf.tanh(l2_out)
    output, out_next_state = self.out(l2_out, out_prev_state)

    # Output state of LSTMNetwork contains the output states of inner modules.
    full_output_state = (l1_next_state, l2_next_state, out_next_state)

    return tf.sigmoid(output), full_output_state

最后，如果您使用的是 Eager 模式，我强烈建议您看看 Sonnet 2 - 这是对 TF 2 / Eager 模式的完全重写。它不向后兼容，但所有相同类型的模块组合都是可能的。 Sonnet 1 主要是为 Graph 模式 TF 编写的，虽然它确实适用于 Eager 模式，但您可能会遇到一些不太方便的事情。

我们与 TensorFlow 团队密切合作，以确保 TF 2 和 Sonnet 2 能够很好地协同工作，因此请查看：(https://github.com/deepmind/sonnet/tree/v2)。 Sonnet 2 应该被认为是 alpha 版本，并且正在积极开发中，所以我们还没有大量示例，但在不久的将来会添加更多示例。