How do I re-batch a tensor in Tensorflow?

Question

I have a model that passes data through a CNN and into an LSTM. But before the LSTM stage I have to reorganise the output of the CNN, as follows:

Input: (600, 512, 1) --(CNN)--> (600, 32)

This (600, 32) tensor is modified by two parameters, k (window size) and s (step size), both integers, which govern the shape of the following tensor (I will call the process 're-batching'):

Re-batch: (600, 32) --> (146, 20, 32)

These dimensions are obtained from k and s using the receptive field equation:

Here n_in=600, n_out evaluates to 146, k=20, p=0, s=4.

This new tensor can then be input into the LSTM, as the shape is (batch_size, timesteps, features), and so the model continues from there...

My problem is that I don't know how to 're-batch' this tensor without breaking the backwards pass of my model. I have tried creating an empty array - np.empty(shape_of_lstm_input) - and iteratively filling each element of the batch axis with the required data, but converting to a numpy array causes a loss of information when trying to minimise the loss, as shown in the following warning message:

WARNING:tensorflow:Gradients do not exist for variables ['list of all parameters in my CNN'] when minimizing the loss.

So instead I tried to follow the same process, but instead using tf.zeros(shape_of_lstm_input) as I thought that would solve the problem of missing gradients, but I instead get this error message:

TypeError: 'tensorflow.python.framework.ops.EagerTensor' object does not support item assignment

Does anyone know if a solution to this problem exists? I don't know for sure how a backwards pass will deal with this 're-batching' of tensors in the forwards pass.

user_1782792 · Answer

Seems like you are trying to create a sliding window out of your data. This is one simple way to get that:

import tensorflow as tf

def sliding_window(x, window_size, stride, axis=0):
    n_in = tf.shape(x)[axis]
    n_out = (n_in - window_size) // stride + 1
    # Just in case n_in < window_size
    n_out = tf.math.maximum(n_out, 0)
    r = tf.expand_dims(tf.range(n_out), 1)
    idx = r * stride + tf.range(window_size)
    return tf.gather(x, idx, axis=axis)

# Test
x = tf.reshape(tf.range(30), [10, 3])
print(x.numpy())
# [[ 0  1  2]
#  [ 3  4  5]
#  [ 6  7  8]
#  [ 9 10 11]
#  [12 13 14]
#  [15 16 17]
#  [18 19 20]
#  [21 22 23]
#  [24 25 26]
#  [27 28 29]]
y = sliding_window(x, 4, 2)
print(y.numpy())
# [[[ 0  1  2]
#   [ 3  4  5]
#   [ 6  7  8]
#   [ 9 10 11]]
# 
#  [[ 6  7  8]
#   [ 9 10 11]
#   [12 13 14]
#   [15 16 17]]
# 
#  [[12 13 14]
#   [15 16 17]
#   [18 19 20]
#   [21 22 23]]
# 
#  [[18 19 20]
#   [21 22 23]
#   [24 25 26]
#   [27 28 29]]]

How do I re-batch a tensor in Tensorflow?

1 Answers