1
$\begingroup$

I have been stuck with a problem like this for a while now. I have an AWS setup with 500 GB of RAM and about 7 GPUs. Now the issue is that each time I try to run my Keras with TensorFlow as back-end code, it runs out of memory. I have found out the reason for this as well. The reason is that each GPU just has 12GB of memory, whereas my model needs more than that. So, how can I run the model such that it uses the memory of all the GPUs combined to load the model and not just rely on the memory of one GPU for loading the entire model and running out of memory? I have tried model parallelism with Keras and it seems to be set-up correctly as on printing the layers , each layer is assigned to the programmed GPU but the model is still trying to load into a single GPU's memory, i.e., just 11GB and soon runs out of memory.

Any idea what's going on?

with tf.device('/gpu:0'):
    x = conv2d_bn(img_input, 32, 3, 3, strides=(2, 2), padding='valid')
    x = conv2d_bn(x, 32, 3, 3, padding='valid')
    x = conv2d_bn(x, 64, 3, 3)
    x = MaxPooling2D((3, 3), strides=(2, 2))(x)

    x = conv2d_bn(x, 80, 1, 1, padding='valid')
    x = conv2d_bn(x, 192, 3, 3, padding='valid')
    x = MaxPooling2D((3, 3), strides=(2, 2))(x)

    # mixed 0, 1, 2: 35 x 35 x 256
    branch1x1 = conv2d_bn(x, 64, 1, 1)

    branch5x5 = conv2d_bn(x, 48, 1, 1)
    branch5x5 = conv2d_bn(branch5x5, 64, 5, 5)

    branch3x3dbl = conv2d_bn(x, 64, 1, 1)
    branch3x3dbl = conv2d_bn(branch3x3dbl, 96, 3, 3)
    branch3x3dbl = conv2d_bn(branch3x3dbl, 96, 3, 3)

    branch_pool = AveragePooling2D((3, 3), strides=(1, 1), padding='same')(x)
    branch_pool = conv2d_bn(branch_pool, 32, 1, 1)
    x = layers.concatenate(
        [branch1x1, branch5x5, branch3x3dbl, branch_pool],
        axis=channel_axis,
        name='mixed0')
    print(x)

with tf.device('/gpu:1'):
    # mixed 1: 35 x 35 x 256
    branch1x1 = conv2d_bn(x, 64, 1, 1)

    branch5x5 = conv2d_bn(x, 48, 1, 1)
    branch5x5 = conv2d_bn(branch5x5, 64, 5, 5)

    branch3x3dbl = conv2d_bn(x, 64, 1, 1)
    branch3x3dbl = conv2d_bn(branch3x3dbl, 96, 3, 3)
    branch3x3dbl = conv2d_bn(branch3x3dbl, 96, 3, 3)

    branch_pool = AveragePooling2D((3, 3), strides=(1, 1), padding='same')(x)
    branch_pool = conv2d_bn(branch_pool, 64, 1, 1)
    x = layers.concatenate(
        [branch1x1, branch5x5, branch3x3dbl, branch_pool],
        axis=channel_axis,
        name='mixed1')

    # mixed 2: 35 x 35 x 256
    branch1x1 = conv2d_bn(x, 64, 1, 1)

    branch5x5 = conv2d_bn(x, 48, 1, 1)
    branch5x5 = conv2d_bn(branch5x5, 64, 5, 5)

    branch3x3dbl = conv2d_bn(x, 64, 1, 1)
    branch3x3dbl = conv2d_bn(branch3x3dbl, 96, 3, 3)
    branch3x3dbl = conv2d_bn(branch3x3dbl, 96, 3, 3)

    branch_pool = AveragePooling2D((3, 3), strides=(1, 1), padding='same')(x)
    branch_pool = conv2d_bn(branch_pool, 64, 1, 1)
    x = layers.concatenate(
        [branch1x1, branch5x5, branch3x3dbl, branch_pool],
        axis=channel_axis,
        name='mixed2')
    print(x)

Edit: Here's the link to the code.

$\endgroup$
2
  • $\begingroup$ You solved your problem @Rueben_v1 ?? $\endgroup$
    – Mike
    Feb 27, 2019 at 15:23
  • $\begingroup$ Not really @Mike , the last time I had checked, model parallelism(in Keras) has not reached the stage where the same model can be distributed across GPUs. I was forced to use another instance with a higher GPU capacity. That being said, you might want to check on the status now. Should you have found a solution, kindly do post it here :) $\endgroup$
    – Reuben_v1
    Mar 9, 2019 at 15:10

1 Answer 1

0
$\begingroup$

What you are describing is model parallelism - different parts of a single model run on different devices, processing a single batch of data together.

There are many ways doing model parallelism in Keras. You probably want to do single-host, multi-device synchronous training - one machine with several GPUs on it. Each device will run a copy of your model (called a replica).

The code looks something like this:

import tensorflow as tf

strategy = tf.distribute.MirroredStrategy()

# Open a strategy scope.
with strategy.scope():
  # Everything that creates variables should be under the strategy scope.
  # In general this is only model construction & `compile()`.
  model = Model(...)
  model.compile(...)

# Train the model on all available devices.
model.fit(train_dataset, validation_data=val_dataset, ...)

# Test the model on all available devices.
model.evaluate(test_dataset)
$\endgroup$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.