How to reconstruct the decoder from an LSTM-AE?

Question

I have a trained LSTM-AE, of which the architecture is as follows:

In brief, I have an LSTM-AE of depth 3, the number of cells on the LSTM layers on the encoder side are [120, 80, 50] (and symmetric for the decoder). I built the model using the code shown on this page. For information, because I want to train the LSTM-AT directly on variable-length time series, so I didn't specify the timestamps in the input layer, which means the model is trained on batches of size 1 (one time series per batch).

I can extract the encoder just fine, but I cannot do the same for the decoder :-(... My goal is to check, given a vector of 50 features (which are extracted by the encoder), whether the decoder can reconstruct the input series.

Here's my attempt so far:

# load the full autoencoder
model = load_model(path_to_model)  

# reconstruct the decoder
in_layer = Input(shape=(None, 50))
time_dist = model.layers[-1]
dec_1 = model.layers[-2]
dec_2 = model.layers[-3]
dec_3 = model.layers[-4]
rep_vec = model.layers[-5]
out_layer = time_dist(dec_1(dec_2(dec_3(rep_vec(in_layer)))))
decoder = Model(in_layer, out_layer, name='decoder')
res = decoder(input_feature)   # input_feature has shape (50,)

I obtained this error:

InvalidArgumentError: slice index 1 of dimension 0 out of bounds. [Op:StridedSlice] name: decoder/repeat/strided_slice/

If you are interested in the full error log...

---------------------------------------------------------------------------
InvalidArgumentError                      Traceback (most recent call last)
Input In [86], in <module>
     13 out_layer = time_dist(dec_1(dec_2(dec_3(rep_vec(in_layer)))))
     14 decoder = Model(in_layer, out_layer, name='decoder')
---> 15 res = decoder(input_feature)

File ~/venv/lib/python3.8/site-packages/tensorflow/python/keras/engine/base_layer.py:1030, in Layer.__call__(self, *args, **kwargs)
   1026   inputs = self._maybe_cast_inputs(inputs, input_list)
   1028 with autocast_variable.enable_auto_cast_variables(
   1029     self._compute_dtype_object):
-> 1030   outputs = call_fn(inputs, *args, **kwargs)
   1032 if self._activity_regularizer:
   1033   self._handle_activity_regularization(inputs, outputs)

File ~/venv/lib/python3.8/site-packages/tensorflow/python/keras/engine/functional.py:420, in Functional.call(self, inputs, training, mask)
    401 @doc_controls.do_not_doc_inheritable
    402 def call(self, inputs, training=None, mask=None):
    403   """Calls the model on new inputs.
    404 
    405   In this case `call` just reapplies
   (...)
    418       a list of tensors if there are more than one outputs.
    419   """
--> 420   return self._run_internal_graph(
    421       inputs, training=training, mask=mask)

File ~/venv/lib/python3.8/site-packages/tensorflow/python/keras/engine/functional.py:556, in Functional._run_internal_graph(self, inputs, training, mask)
    553   continue  # Node is not computable, try skipping.
    555 args, kwargs = node.map_arguments(tensor_dict)
--> 556 outputs = node.layer(*args, **kwargs)
    558 # Update tensor_dict.
    559 for x_id, y in zip(node.flat_output_ids, nest.flatten(outputs)):

File ~/venv/lib/python3.8/site-packages/tensorflow/python/keras/engine/base_layer.py:1030, in Layer.__call__(self, *args, **kwargs)
   1026   inputs = self._maybe_cast_inputs(inputs, input_list)
   1028 with autocast_variable.enable_auto_cast_variables(
   1029     self._compute_dtype_object):
-> 1030   outputs = call_fn(inputs, *args, **kwargs)
   1032 if self._activity_regularizer:
   1033   self._handle_activity_regularization(inputs, outputs)

File ~/venv/lib/python3.8/site-packages/tensorflow/python/keras/layers/core.py:919, in Lambda.call(self, inputs, mask, training)
    915   return var
    917 with backprop.GradientTape(watch_accessed_variables=True) as tape,\
    918     variable_scope.variable_creator_scope(_variable_creator):
--> 919   result = self.function(inputs, **kwargs)
    920 self._check_variables(created_variables, tape.watched_variables())
    921 return result

File D:/PhD/Code/feature_learning/train_models/train_lstmae.py:30, in repeat_vector(args)

File ~/venv/lib/python3.8/site-packages/tensorflow/python/util/dispatch.py:206, in add_dispatch_support.<locals>.wrapper(*args, **kwargs)
    204 """Call target, and fall back on dispatchers if there is a TypeError."""
    205 try:
--> 206   return target(*args, **kwargs)
    207 except (TypeError, ValueError):
    208   # Note: convert_to_eager_tensor currently raises a ValueError, not a
    209   # TypeError, when given unexpected types.  So we need to catch both.
    210   result = dispatch(wrapper, args, kwargs)

File ~/venv/lib/python3.8/site-packages/tensorflow/python/ops/array_ops.py:1040, in _slice_helper(tensor, slice_spec, var)
   1038   var_empty = constant([], dtype=dtypes.int32)
   1039   packed_begin = packed_end = packed_strides = var_empty
-> 1040 return strided_slice(
   1041     tensor,
   1042     packed_begin,
   1043     packed_end,
   1044     packed_strides,
   1045     begin_mask=begin_mask,
   1046     end_mask=end_mask,
   1047     shrink_axis_mask=shrink_axis_mask,
   1048     new_axis_mask=new_axis_mask,
   1049     ellipsis_mask=ellipsis_mask,
   1050     var=var,
   1051     name=name)

File ~/venv/lib/python3.8/site-packages/tensorflow/python/util/dispatch.py:206, in add_dispatch_support.<locals>.wrapper(*args, **kwargs)
    204 """Call target, and fall back on dispatchers if there is a TypeError."""
    205 try:
--> 206   return target(*args, **kwargs)
    207 except (TypeError, ValueError):
    208   # Note: convert_to_eager_tensor currently raises a ValueError, not a
    209   # TypeError, when given unexpected types.  So we need to catch both.
    210   result = dispatch(wrapper, args, kwargs)

File ~/venv/lib/python3.8/site-packages/tensorflow/python/ops/array_ops.py:1213, in strided_slice(input_, begin, end, strides, begin_mask, end_mask, ellipsis_mask, new_axis_mask, shrink_axis_mask, var, name)
   1210 if strides is None:
   1211   strides = ones_like(begin)
-> 1213 op = gen_array_ops.strided_slice(
   1214     input=input_,
   1215     begin=begin,
   1216     end=end,
   1217     strides=strides,
   1218     name=name,
   1219     begin_mask=begin_mask,
   1220     end_mask=end_mask,
   1221     ellipsis_mask=ellipsis_mask,
   1222     new_axis_mask=new_axis_mask,
   1223     shrink_axis_mask=shrink_axis_mask)
   1225 parent_name = name
   1227 if var is not None:

File ~/venv/lib/python3.8/site-packages/tensorflow/python/ops/gen_array_ops.py:10505, in strided_slice(input, begin, end, strides, begin_mask, end_mask, ellipsis_mask, new_axis_mask, shrink_axis_mask, name)
  10503   return _result
  10504 except _core._NotOkStatusException as e:
> 10505   _ops.raise_from_not_ok_status(e, name)
  10506 except _core._FallbackException:
  10507   pass

File ~/venv/lib/python3.8/site-packages/tensorflow/python/framework/ops.py:6897, in raise_from_not_ok_status(e, name)
   6895 message = e.message + (" name: " + name if name is not None else "")
   6896 # pylint: disable=protected-access
-> 6897 six.raise_from(core._status_to_exception(e.code, message), None)

File <string>:3, in raise_from(value, from_value)

InvalidArgumentError: slice index 1 of dimension 0 out of bounds. [Op:StridedSlice] name: decoder/repeat/strided_slice/

I appreciate very much any advice you would give me!

Edit

Here is the code I used to build the mode:

import tensorflow as tf
from tensorflow.keras.layers import *
from tensorflow.keras.initializers import GlorotUniform
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.models import Sequential, Model
from tensorflow.keras.backend import shape

def repeat_vector(args):
    """Builds the repeat vector layer dynamically by the size of the input series"""
    layer_to_repeat = args[0]
    sequence_layer = args[1]
    return RepeatVector(shape(sequence_layer)[1])(layer_to_repeat)

n_atts = 3   # time series of 3 measurements
n_units = [120, 80, 50]   # encoder - 1st layer: 120, 2nd layer: 80, 3rd layer: 50 (and symmetric for decoder)
n_layers = len(n_units)
init = GlorotUniform(seed=420)
reg = None
optimizer = Adam(learning_rate=0.0001)
activ = 'tanh'
loss_metric = 'mse'

inputs = Input(shape=(None, n_atts), name='input_layer')

# the encoder
encoded = LSTM(n_units[0], name='encoder_1', return_sequences=(n_layers != 1), kernel_initializer=init,
               kernel_regularizer=reg, activation=activ)(inputs)
for i in range(1, n_layers):
    if i != n_layers - 1:
        encoded = LSTM(n_units[i], name='encoder_{}'.format(i + 1), return_sequences=(n_layers != 1),
                       kernel_initializer=init, kernel_regularizer=reg, activation=activ)(encoded)
    else:
        encoded = LSTM(n_units[i], name='encoder_{}'.format(i + 1), return_sequences=False,
                       kernel_initializer=init, kernel_regularizer=reg, activation=activ)(encoded)

# repeat the vector (plug the encoder to the decoder)
repeated = Lambda(repeat_vector, output_shape=(None, n_units[-1]), name='repeat')([encoded, inputs])

# the decoder
decoded = LSTM(n_units[n_layers - 1], return_sequences=True, name='decoder_1',
               kernel_initializer=init, kernel_regularizer=reg, activation=activ)(repeated)  # first layer
for i in range(1, n_layers):
    decoded = LSTM(n_units[n_layers - 1 - i], return_sequences=True, name='decoder_{}'.format(i + 1),
                   kernel_initializer=init, kernel_regularizer=reg, activation=activ)(decoded)

# last layer
tdist = TimeDistributed(Dense(n_atts))(decoded)

# compile the model
model = Model(inputs, tdist, name='lstm-ae')
model.compile(optimizer=optimizer, loss=loss_metric)

For information, I use tensorflow 2.5.

Because the number of units is read from a config file, I wrote the code this way to add the layers programmatically.