tensorflow dataset shuffle examples instead of batches - tensorflow

Question

How do I get a tensorflow dataset in batch mode to shuffle across all the samples? It is only shuffling the batches.
Below is a program that makes a dataset of 1000 items and goes through 10 epochs of it in batches of 5. I have shuffle() turned on. I can see that tensorflow groups the dataset into 200 batches of 5 examples each, and the shuffle is across those batches. I want each new batch to be a random sample of the original 1000 examples, not a sample of the 200 original batches.
That is, this program:
import numpy as np
import tensorflow as tf
import random
def rec2tfrec_example(rec):
def _int64_feat(value):
arr_value = np.empty([1], dtype=np.int64)
arr_value[0] = value
return tf.train.Feature(int64_list=tf.train.Int64List(value=arr_value))
feat = {
'uid': _int64_feat(rec['uid']),
}
return tf.train.Example(features=tf.train.Features(feature=feat)).SerializeToString()
def parse_example(tfrec_serialized_string):
feat = {
'uid': tf.FixedLenFeature([], tf.int64),
}
return tf.parse_example(tfrec_serialized_string, feat)
def write_tfrecs_to_file(fname, recs):
recwriter = tf.python_io.TFRecordWriter(fname)
for rec in recs:
recwriter.write(bytes(rec))
recwriter.close()
def check_shuffle(sess, tfrec_output_filename, data, N, batch_size):
epochs = 10
dataset = tf.data.TFRecordDataset(tfrec_output_filename) \
.batch(batch_size) \
.repeat(epochs) \
.shuffle(2*N) \
.map(parse_example, num_parallel_calls=2)
tf_iter = dataset.make_initializable_iterator()
get_next = tf_iter.get_next()
sess.run(tf_iter.initializer)
num_batches = N//batch_size
for epoch in range(epochs ):
for batch in range(N//batch_size):
tfres = sess.run(get_next)
print("epoch=%4d batch=%d uid=%s" % (epoch, batch, tfres['uid']))
def main(N=1000, batch_size=5, tfrec_output_filename='tfrec_testing.tfrecords'):
tf.reset_default_graph()
data = [{'uid': uid } for uid in range(N)]
tfrec_strings = [rec2tfrec_example(rec) for rec in data]
write_tfrecs_to_file(tfrec_output_filename, tfrec_strings)
with tf.Session() as sess:
check_shuffle(sess, tfrec_output_filename, data, N, batch_size)
if __name__ == '__main__':
main()
produces output like:
epoch= 9 batch=186 uid=[685 686 687 688 689]
epoch= 9 batch=187 uid=[235 236 237 238 239]
epoch= 9 batch=188 uid=[520 521 522 523 524]
epoch= 9 batch=189 uid=[135 136 137 138 139]
epoch= 9 batch=190 uid=[95 96 97 98 99]
epoch= 9 batch=191 uid=[290 291 292 293 294]
epoch= 9 batch=192 uid=[230 231 232 233 234]
epoch= 9 batch=193 uid=[215 216 217 218 219]

Answer 1

ah, the order of batch and shuffle matters, if I set up the dataset like
dataset = tf.data.TFRecordDataset(tfrec_output_filename) \
.shuffle(2*N) \
.batch(batch_size) \
.repeat(epochs) \
.map(parse_example, num_parallel_calls=2)
with shuffle before batch, then it works.