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

我通过子类化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.

问题在于，上述代码既不是顺序模型，也不是函数模型。我试图在第一个掩蔽层中指定批处理输入形状，但它不起作用

---

如何解决上述错误？我目前使用的是tensorflow 2.0 rc0。

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

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

---

你找到解决办法了吗？我也有同样的问题。非常令人沮丧。