convLSTM2d w/ functional API - tensorflow

Question

I have an autoencoder for image compression, where the encoded tensor has the shape: (batch_size, 12, 64, 48).
batch_size is the number of images being fed in a batch,
12 is the number of channels of this last encoder layer, which has a
64x48 width/height.
I want to input this to a ConvLSTM2D layer, and i would like the output of the ConvLSTM2D to have the same dimension as the input of the ConvLSTM2D.
The intention is to see image reconstruction on a video sequence, rather than unordered images from a dataset.
Placing a ConvLSTM2d between an encoder/decoder in a autoencoder architecture has been difficult, especially because most examples use the Sequential API, and i want to use the functional API in Keras.
I tried reshaping the input but the error persists
import tensorflow as tf
import tensorflow.keras.backend as K
def LSTM_layer(input):
input = tf.keras.backend.expand_dims(input, axis=-1)
lstm1 = tf.keras.layers.ConvLSTM2D(filters=12, kernel_size=(3, 3), strides=(1, 1), data_format="channels_first",
input_shape=(None, 12, 64, 48),
padding='same', return_sequences=True)(input)
return lstm1
def build_model(input_shape):
#create an input with input_shape as the size
input_ = tf.keras.Input(shape=input_shape, name="input_node")
lstm_features = LSTM_layer(input_)
model = tf.keras.Model(inputs=input_, outputs=[lstm_features])
return model
def main():
input_shape = (12, 64, 48) #this is the size of the tensor which is outputted by my encoder, with channels_first assumed
model = build_model(input_shape)
if __name__ == '__main__':
main()
Unfortunately, this is throwing this error:
Traceback (most recent call last):
File "lstm.py", line 29, in <module>
main()
File "lstm.py", line 26, in main
model = build_model(input_shape)
File "lstm.py", line 20, in build_model
model = tf.keras.Model(inputs=input_, outputs=[lstm_features])
File "/home/hallab/.local/lib/python3.5/site-packages/tensorflow/python/keras/engine/training.py", line 121, in __init__
super(Model, self).__init__(*args, **kwargs)
File "/home/hallab/.local/lib/python3.5/site-packages/tensorflow/python/keras/engine/network.py", line 80, in __init__
self._init_graph_network(*args, **kwargs)
File "/home/hallab/.local/lib/python3.5/site-packages/tensorflow/python/training/checkpointable/base.py", line 474, in _method_wrapper
method(self, *args, **kwargs)
File "/home/hallab/.local/lib/python3.5/site-packages/tensorflow/python/keras/engine/network.py", line 224, in _init_graph_network
'(thus holding past layer metadata). Found: ' + str(x))
ValueError: Output tensors to a Model must be the output of a TensorFlow `Layer` (thus holding past layer metadata). Found: Tensor("conv_lst_m2d/transpose_1:0", shape=(?, 12, 12, 48, 1), dtype=float32)
Most posts about this error instruct to wrap the operation in a lambda.. but i am not implementing a custom operation here, this should be a keras tf layer... right?
Also, in my implementation, i want the output tensor from the LSTM unit to be the same as the input, can i get some feedback about that as well?
Thank you.

Answer 1

You could use Lambda to wrap the output form K.expand_dims before input it to next layer like this:
import tensorflow as tf
import tensorflow.keras.backend as K
from tensorflow.keras.layers import Lambda
def expand_dims(x):
return K.expand_dims(x, 1)
def expand_dims_output_shape(input_shape):
return (input_shape[0], 1, input_shape[1])
def LSTM_layer(input_):
lstm1 = Lambda(expand_dims, expand_dims_output_shape)(input_)
lstm1 = tf.keras.layers.ConvLSTM2D(filters=12, kernel_size=(3, 3), strides=(1, 1), data_format="channels_first", padding='same', return_sequences=False)(lstm1)
return lstm1