Cfrgtkky

I have a working solution for my problem, but I find it relatively slow, and I would be surprised if there is no better way of doing what I want. See the description of the problem below.

CURRENT SOLUTION: (comes first, as it might be simpler for experienced users of tensorflow to infer what I want to do from this code snippet)

def stack_batch_dim(*x):

    return {key: tf.stack([b[key] for b in x], axis=1) for key in x[0]}





dataset = tf.data.TFRecordDataset(filenames)

dataset = dataset.map(map_func=parse_box2d_protobuf, num_parallel_calls=8)

dataset = dataset.repeat()



# place the "cursors":

datasets = [dataset.skip(t * inter_cursor_space) for t in range(batch_size)]

dataset = tf.data.Dataset.zip(tuple(datasets))



# consume time_depth samples at a time

dataset = dataset.batch(time_depth)



# the actual batch dimension is a tuple and not a tensor

# need to "stack" the tuple into a tensor

dataset = dataset.map(stack_batch_dim)

iterator = dataset.make_initializable_iterator()

So I would like train an LSTM network with a sequence of data points. The data I have is stored in tfrecords, (the recommended way of storing data with tensorflow).

dataset = tf.data.TFRecordDataset(filenames)

dataset = dataset.map(map_func=parse_box2d_protobuf, num_parallel_calls=8)



# parse_box2d_protobuf is a function that transforms raw bytes into

# a python structure of the form:

# {"positions": tensor_of_shape_4,

#  "speeds": tensor_of_shape_4,

#  "action": tensor_of_shape_4,

#  "vision": tensor_of_shape_h_w_3,

#  "index": tensor_int64}





iterator = dataset.make_initializable_iterator()

next_sample = iterator.get_next()

with tf.Session() as sess:

    sess.run(iterator.initializer)

    print(sess.run(next_sample))

    # prints a dictionary as expected.

Now what I want to do is to read that sequence in order to feed it in an LSTM in batches.
As you might expect, the "index" key in the dictionary corresponds to an integer that is the index of the data sample in the sequence (starting from 0):

with tf.Session() as sess:

    sess.run(iterator.initializer)

    for i in range(10, 15):

        print(sess.run(next_sample["index"]))

        # prints successively 0, 1, 2, 3 and 4

Now let's say I want a batch size of 4, and I want to unfold the LSTM net for 3 iterations. This means I need to read the sequence at 4 positions simultaneously (4 cursors), each cursor reading samples 3 by 3. Let's add a parameter that controls the space between the cursors:

batch_size = 4

time_depth = 3

inter_cursor_space = 5

How can I define an iterator such that calling iterator.get_next()["index"] in a session returns

with tf.Session() as sess:

    sess.run(iterator.initializer)



    # first call:

    print(sess.run(next_sample["index"]))



    # should print

    # [[0, 5, 10, 15],

    #  [1, 6, 11, 16],

    #  [2, 7, 12, 17]]



    # second call:

    print(sess.run(next_sample["index"]))



    # should print

    # [[3, 8, 13, 18],

    #  [4, 9, 14, 19],

    #  [5, 10, 15, 20]]



    # shape = [time_depth, batch_size]

Such the this iterator is compatible with the tf.nn.static_rnn function.