tf.data.Iterator with a multi GPU setup

Question

I have looked at the cifar10 multi-GPU implementation to draw inspiration for parallelizing my own GPU trained model.

My model consumes data from TFRecords, which are iterated through the tf.data.Iterator class. So given 2 GPUs what I am trying to do is call iterator.get_next() on the CPU one time for each GPU (twice for example) do some preprocessing ,embedding lookup and other CPU related stuff and then feed the two batches into the GPUs.

Pseudo code:

with tf.device('/cpu:0'):
    batches = []
    for gpu in multiple_gpus:
        single_gpu_batch = cpu_function(iterator.get_next())
        batches.append(single_gpu_batch)

    ....................

for gpu, batch in zip(multiple_gpus, batches):
    with tf.device('/device:GPU:{}'.format(gpu.id):
        single_gpu_loss = inference_and_loss(batch)
        tower_losses.append(single_gpu_loss)
        ...........
        ...........

total_loss = average_loss(tower_losses)

The problem is, that if there is only 1 or less examples to be drawn from the data and I call iterator.get_next() twice a tf.errors.OutOfRange exception will be raised and the data of the first call of iterator.get_next() (which actually didn't fail, only the second one) will never be passed through the GPU.

I thought about drawing the data in one iterator.get_next() call and splitting it later, but tf.split fails of the batch size is not dividable by the number of GPUs.

What is the right way to implement consuming from iterator in a multi-GPU setup?

user_9934740 · Answer

I think the second suggestion is the easiest way to go. In order to avoid the splitting problem on the last batch, you can use the drop_remainder option in dataset.batch; Or if you need to see all data, then one possible solution is to explicitly set the dimensions based on the size of the drawn batch, so that the splitting operation never fails:

dataset = dataset.batch(batch_size * multiple_gpus)
iterator = dataset.make_one_shot_iterator()
batches = iterator.get_next()

split_dims = [0] * multiple_gpus
drawn_batch_size = tf.shape(batches)[0]

Either in a greedy manner, i.e., fits batch_size tensors on each device until runs out

#### Solution 1 [Greedy]: 
for i in range(multiple_gpus):
  split_dims[i] = tf.maximum(0, tf.minimum(batch_size, drawn_batch_size))
  drawn_batch_size -= batch_size

or in a more spread manner to ensure that each device get at least one sample (assuming multiple_gpus < drawn_batch_size)

### Solution 2 [Spread]
drawn_batch_size -= - multiple_gpus
for i in range(multiple_gpus):
  split_dims[i] = tf.maximum(0, tf.minimum(batch_size - 1, drawn_batch_size)) + 1
  drawn_batch_size -= batch_size

## Split batches
batches = tf.split(batches, split_dims)

tf.data.Iterator with a multi GPU setup

1 Answers