tf。数据具有多GPU设置的迭代器

我已经看过了 cifar10 multi-GPU implementation 为我自己的GPU训练模型的并行化提供灵感。

我的模型使用TFRecords中的数据,这些数据通过 tf.data.Iterator 班因此,给定2 GPU,我要做的是调用 iterator.get_next() 在CPU上,为每个GPU执行一次(例如两次)预处理、嵌入查找和其他与CPU相关的操作,然后将这两个批次输入GPU。

伪代码:

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)

问题是,如果从数据中只提取了1个或更少的示例,我会调用 迭代器。下一个 一天两次 tf.errors.OutOfRange 将引发异常,并且

我想把数据画在一张纸上 迭代器。下一个 tf.split 批处理大小的大小不能除以GPU的数量。