Out Of Memory when running multi-gpu cnn with TensorFlow - tensorflow

Question

I'm trying to run a simple cnn on cifar10, combining code from 2 examples:
https://github.com/aymericdamien/TensorFlow-Examples/blob/master/examples/6_MultiGPU/multigpu_cnn.py
https://github.com/exelban/tensorflow-cifar-10
I'm getting OOM errors.
I first tried the code with the complete cnn , without multi-gpu support, and it is working ok. Next I used the multi-gpu code, ran ok too.
Combining them is not working.
with tf.device('/cpu:0'):
tower_grads = []
reuse_vars = False
# tf Graph input
X = tf.placeholder(tf.float32, shape=[None, _IMAGE_SIZE * _IMAGE_SIZE * _IMAGE_CHANNELS], name='Input')
Y = tf.placeholder(tf.float32, shape=[None, _NUM_CLASSES], name='Output')
phase = tf.placeholder(tf.bool, name='phase')
# learning_rate = tf.placeholder(tf.float32, shape=[], name='learning_rate')
keep_prob = tf.placeholder(tf.float32)
global_step = tf.get_variable(name='global_step', trainable=False, initializer=0)
# Loop over all GPUs and construct their own computation graph
for i in range(_NUM_GPUS):
with tf.device('/gpu:{}'.format(i)):
# learning_rate = tf.placeholder(tf.float32, shape=[], name='learning_rate')
# keep_prob = tf.placeholder(tf.float32)
# Split data between GPUs
_x = X[i * _BATCH_SIZE: (i+1) * _BATCH_SIZE]
_y = Y[i * _BATCH_SIZE: (i+1) * _BATCH_SIZE]
print("x shape:",_x.shape)
print("y shape:",_y.shape)
# Because Dropout have different behavior at training and prediction time, we
# need to create 2 distinct computation graphs that share the same weights.
_x = tf.reshape(_x, [-1, _IMAGE_SIZE, _IMAGE_SIZE, _IMAGE_CHANNELS], name='images')
# Create a graph for training
logits_train, y_pred_cls = feed_net(_x, _NUM_CLASSES, keep_prob, reuse=reuse_vars, is_training=True)
# Create another graph for testing that reuse the same weights
logits_test, y_pred_cls = feed_net(_x, _NUM_CLASSES, keep_prob, reuse=True, is_training=False)
# Define loss and optimizer (with train logits, for dropout to take effect)
loss_op = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=logits_train, labels=_y))
optimizer = tf.train.AdamOptimizer(learning_rate=0.001)
grads = optimizer.compute_gradients(loss_op)
# Only first GPU compute accuracy
if i == 0:
# Evaluate model (with test logits, for dropout to be disabled)
correct_pred = tf.equal(tf.argmax(logits_test, 1), tf.argmax(_y, 1))
accuracy = tf.reduce_mean(tf.cast(correct_pred, tf.float32))
reuse_vars = True
tower_grads.append(grads)
tower_grads = average_gradients(tower_grads)
train_op = optimizer.apply_gradients(tower_grads)
The error is happening when running with more than 1 gpu (got 4), after about 90 iterations (less than one epoch)
ResourceExhaustedError: Ran out of GPU memory when allocating 0 bytes for
[[Node: softmax_cross_entropy_with_logits_sg_3 = SoftmaxCrossEntropyWithLogits[T=DT_FLOAT, _device="/job:localhost/replica:0/task:0/device:GPU:3"](softmax_cross_entropy_with_logits_sg_3/Reshape, softmax_cross_entropy_with_logits_sg_3/Reshape_1)]]
Hint: If you want to see a list of allocated tensors when OOM happens, add report_tensor_allocations_upon_oom to RunOptions for current allocation info.
[[Node: main_params/map/while/Less_1/_206 = _Send[T=DT_BOOL, client_terminated=false, recv_device="/job:localhost/replica:0/task:0/device:GPU:0", send_device="/job:localhost/replica:0/task:0/device:GPU:0", send_device_incarnation=1, tensor_name="edge_1905_main_params/map/while/Less_1", _device="/job:localhost/replica:0/task:0/device:GPU:0"](main_params/map/while/Less_1)]]
Hint: If you want to see a list of allocated tensors when OOM happens, add report_tensor_allocations_upon_oom to RunOptions for current allocation info.
UPDATE:
The problem was with how the data was divided across the GPUs.
I used tf.split(X, _NUM_GPUS) for the data and the labels, then I could assign each GPU with it's right data chunk.

Answer 1

Here's the solution:
The problem was with how the data was divided across the GPUs.
I used tf.split(X, _NUM_GPUS) for the data and the labels, then I could assign each GPU with it's right data chunk. Also , only one GPU is running accuracy so it needed to get the full sized data.