Rllib TensorFlow自定义LSTM模型提供InvalidArgumentError:不兼容的形状错误(在LSTM层中)

时间:2020-09-30 20:50:00

标签: python tensorflow keras ray rllib

我一直在尝试使用RLLib建立自定义LSTM模型,但是由于某些原因,在尝试训练时,我的LSTM层出现了形状不兼容的错误。特别是,此错误似乎与批次大小有关,因为针对不兼容形状列出的尺寸随我的批次大小值线性变化。我的模型代码如下:

class CustomLSTMModel(RecurrentNetwork):
    """Example of using the Keras functional API to define a RNN model."""

    def __init__(
        self,
        obs_space,
        action_space,
        num_outputs,
        model_config,
        name,
        hiddens_size=64,
        cell_size=64,
    ):
        super(CustomLSTMModel, self).__init__(
            obs_space, action_space, num_outputs, model_config, name
        )
        self.cell_size = cell_size

        # Define input layers
        input_layer = tf.keras.layers.Input(
            shape=(obs_space.shape[0], obs_space.shape[1]), name="inputs"
        )
        state_in_h = tf.keras.layers.Input(shape=(cell_size,), name="h")
        state_in_c = tf.keras.layers.Input(shape=(cell_size,), name="c")

        # FC layer
        dense1 = tf.keras.layers.Dense(hiddens_size, name="dense_1")(input_layer)
  
        # LSTM layer
        lstm_1_out, state_1_h, state_1_c = tf.keras.layers.LSTM(
            cell_size,
            return_state=True,
            name="lstm_1",
        )(
            inputs=dense1,
            initial_state=[state_in_h, state_in_c],
        )

        # Postprocess 
        logits = tf.keras.layers.Dense(
            self.num_outputs,
            activation=None,
            name="logits",
        )(lstm_1_out)
        values = tf.keras.layers.Dense(
            1,
            activation=None,
            name="values",
        )(lstm_1_out)

        # Create the RNN model
        self.rnn_model = tf.keras.Model(
            inputs=[input_layer, state_in_h, state_in_c],
            outputs=[logits, values, state_1_h, state_1_c],
        )
        self.register_variables(self.rnn_model.variables)
        self.rnn_model.summary()

    @override(ModelV2)
    def forward(self, input_dict, state, seq_lens):
        """Custom forward pass for inputs that already have time dimension."""
        inputs = input_dict["obs"]
        output, new_state = self.forward_rnn(inputs, state, seq_lens)
        return tf.reshape(output, [-1, self.num_outputs]), new_state

    @override(RecurrentNetwork)
    def forward_rnn(self, inputs, state, seq_lens):
        model_out, self._value_out, h, c = self.rnn_model([inputs] + state)
        return model_out, [h, c]

    @override(ModelV2)
    def get_initial_state(self):
        return [
            np.zeros(self.cell_size, np.float32),
            np.zeros(self.cell_size, np.float32),
        ]

    @override(ModelV2)
    def value_function(self):
        return tf.reshape(self._value_out, [-1])

这非常接近于现有的RNN模型的RLLib代码,位于以下位置:https://github.com/ray-project/ray/blob/8abe13023f72dabeb233278c0b7b4f80f84c5cc3/rllib/models/tf/recurrent_net.py

我必须进行的主要更改是在覆盖的前向传递函数中,因为作为该函数的输入提供的input_dict已经添加了时间维度,因此进入的维度为(?,5, 50)(?,10、50)等。因此,我不需要它们用于添加新时间维度的add_time_dimension函数。我认为的另一个重要变化是如何在模型类中定义输入层的形状。实际上,我得到的错误基本上与此相同,只是形状不匹配的值不同:Dimensions error, in seq2seq model (op: 'Add') with input shapes: [512,64], [5739,64]

任何想法将不胜感激!

0 个答案:

没有答案