Как сделать так, чтобы Tensorflow 2.0 keras LSTM возвращал последовательности во время обучения, но не во время вывода

Question

Предположим последовательность [1,2,3,4,5,6].

При обучении LSTM полезно предварительно обработать эту последовательность так, чтобы x = [1,2,3, 4,5] и у = [2,3,4,5,6]. Затем модель можно научить делать прогнозы для каждого временного шага, а не только для последнего. Частично это достигается установкой return_sequence в True в реализации tf.keras в LSTM. В моей модели последовательности из каждого временного шага затем передаются через плотный слой для классификации на каждом временном шаге.

Во время вывода я хотел бы, чтобы моя модель выводила прогноз только для последнего временного шага. Я хотел бы продолжать использовать обученные веса из моего плотного слоя для выполнения классификации. Это может быть достигнуто установкой return_state в True на уровне LSTM.

Можно также попытаться вместо этого извлечь последний элемент во всей последовательности, но этот подход не срабатывает, когда входной сигнал прерывистый, а некоторые тензоры дополняются средним потоком, особенно если к этим дополненным тензорам применяется маска.

Поэтому я хочу создать собственный слой LSTM, который ведет себя по-разному во время обучения и вывода. Такое переключаемое поведение достижимо в слоях Keras путем включения переменной training в функцию call.

Я попытался выполнить это переключение с помощью вложенного слоя, используя следующий код:

class FlexibleOutputLSTM(tf.keras.layers.Layer):
    def __init__(self, lstm_units):
        super(FlexibleOutputLSTM, self).__init__()
        self.lstm_units = lstm_units

    def build(self, input_shape):
        self.lstm = tf.keras.layers.LSTM(units=self.lstm_units, return_sequences=True, return_state=True)

    def call(self, inputs, training=False):
        whole_sequence_output, final_memory_state, _ = self.lstm(inputs)
        if training:
            return whole_sequence_output
        return final_memory_state

Но я получаю эту ошибку:

/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/framework/ops.py in __iter__(self)
    541     shape = self._shape_tuple()
    542     if shape is None:
--> 543       raise TypeError("Cannot iterate over a tensor with unknown shape.")
    544     if not shape:
    545       raise TypeError("Cannot iterate over a scalar tensor.")

TypeError: Cannot iterate over a tensor with unknown shape.

Я подозреваю, что это потому, что я неправильно обрабатываю методы init и / или build, но не вижу способа передать форму ввода .

Я вижу следующий пост, предлагающий лямбда-слой, но, как уже упоминалось выше, активации для последнего временного шага во всей последовательности для дополненного входного тензора не достаточно. Когда у этого тензора есть маска, активациями для последнего временного шага будут все нули. Как получить только последний вывод модели последовательности в Керасе?