I am trying to use the TensorFlow (v1.13) Dataset API to save and load long sequences for a stateful RNN.
Basically, lets say I have n_seq
sequences, each fixed to a length of 120.
I want to load batches to train my stateful RNN according to this post, basically:
batch 1 batch 2 batch 3 batch 4
element0 s21 s22 s23 s24 ...
element1 s11 s12 s13 s14
where sij
denotes j-th window of the i-th sequence. This is done based on how a stateful RNN works in TF, as when stateful=True
, "the last state for each sample at index i in a batch will be used as initial state for the sample of index i in the following batch." (source)
Initial Solution
I am saving each sequence to its own *.tfrecord file. Then I am loading them using:
tf_dataset = tf.data.Dataset.list_files(path, shuffle=False, seed=1)
tf_dataset = tf_dataset.interleave(
lambda filename: tf.data.TFRecordDataset(filename).map(parse_func),
cycle_length=BATCH_SIZE,
block_length=1
)
tf_dataset = tf_dataset.batch(batchsize)
This actually mostly works EXCEPT for the final files. Namely if n_files % BATCH_SIZE != 0
, the finale files will be together in the same batch.
Is there a better way to handle this?
NOTE: This question was previously asked on stackoverflow.coom, but deleted due to lack of response. I think this will be a better home for it.