This is my code for generating augmented training batches in Tensorflow

# x_image is the input image tensor
# y_true is the input onehot tensor

# The input tensors to the batch generator. 
# These will be fed with the entire training set
images_tensor = tf.placeholder(tf.float32, shape=self.x_image.shape)
onehots_tensor = tf.placeholder(tf.float32, shape=self.y_true.shape)

# Create a dataset object
data = tf.data.Dataset.from_tensor_slices((images_tensor, onehots_tensor))

# Augment the images
data = data.map(lambda x, y: (self.augment_fn(x), y), num_parallel_calls=32)

# Shuffle them and repeat (images_train is actual data)
data = data.apply(tf.contrib.data.shuffle_and_repeat(len(images_train), None)) 

# Create a batch
data = data.batch(batch_size)

# Pre-fetch a batch worth of data
data = data.prefetch(batch_size)

# Create batch iterator
iterator = data.make_initializable_iterator()
next_element = iterator.get_next()
init_op = iterator.initializer

# Initialise the generator
# images_train and onehot_train are the actual training data (numpy)
session.run(init_op, feed_dict={images_tensor: images_train, onehots_tensor: onehot_train})

Now inside the training loop I get new batches like:

X, y = session.run(next_element)


  • I put the prefetch size to the batch_size. What would be the point of having more or less? Some examples I see have prefetch size only 5 etc.
  • Similarly, I put the randomisation size to that of the whole training set, so that the batches are as random as possible. Is there any case for using less than this? Maybe if memory is a problem and the images are being loading from disk?
  • \$\begingroup\$ Could you add the python version tag? \$\endgroup\$ – t3chb0t Aug 7 '18 at 12:43

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.