子类化keras.model时指定批处理大小

时间:2019-09-09 20:09:50

标签: keras lstm

我通过子类化keras.Model实现了一个包含LSTM层的模型。 以下是我使用的代码。

import tensorflow as tf
from tensorflow import keras

class SequencePredictor(keras.Model):
    def __init__(self, cell_size):
        super(SequencePredictor, self).__init__(self)
        self._mask = keras.layers.Masking(
            mask_value=0.0)
        self._lstm = keras.layers.LSTM(cell_size, return_sequences=True,
                stateful=True)
        self._dense = tf.keras.layers.Dense(1)

    def call(self, inputs, training=False):
        out = self._mask(inputs)
        out = self._lstm(out)
        return self._dense(out)


sequence_predictor = SequencePredictor(cell_size=10)
train_dataset = DataQueue(5000, 500).dataset

sequence_predictor.compile(
        loss="mean_squared_error",
        optimizer=keras.optimizers.Adam(lr=0.05),
        metrics=['mse'])
sequence_predictor.fit(train_dataset, epochs=50)

上面的代码不适用于以下错误消息。

  ValueError: If a RNN is stateful, it needs to know its batch size. Specify the batch size of your input tensors: 
    - If using a Sequential model, specify the batch size by passing a `batch_input_shape` argument to your first layer.
    - If using the functional API, specify the batch size by passing a `batch_shape` argument to your Input layer.

问题在于上述代码既不是顺序模型也不是功能模型。我尝试在第一个Masking层中指定batch_input_shape,但这不起作用。

如何解决以上错误?我目前正在使用tensorflow 2.0 rc0。

1 个答案:

答案 0 :(得分:0)

实例化后,使用所需的输入形状构建LSTM层。

self._lstm = keras.layers.LSTM(cell_size, return_sequences=True, stateful=True)
self._lstm.build((input_shape))