Модель параллелизма не работает? Все графические процессоры не используются?

Question

Я застрял с такой проблемой какое-то время. У меня есть установка AWS с 500 ГБ оперативной памяти и около 7 графических процессоров. Теперь проблема в том, что каждый раз, когда я пытаюсь запустить свои кераты с тензорным потоком в качестве внутреннего кода, ему не хватает памяти. Я также выяснил причину этого. Причина в том, что у каждого вашего GPU просто есть 12 ГБ памяти, а моей модели нужно больше. Итак, как я могу запустить модель так, чтобы она использовала память всех графических процессоров, объединенных для загрузки модели, а не просто полагалась на память одного графического процессора для загрузки всей модели и исчерпания памяти? Я пробовал параллелизм модели с кератами, и, похоже, он настроен правильно, так как при печати слоев каждый слой назначается запрограммированному графическому процессору, но модель все еще пытается загрузить в память одного графического процессора, т.е. только 11 ГБ и скоро заканчивается память.

Есть идеи, что происходит?

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: Вот ссылка на код . Также можно отметить, что в настоящее время я просто передаю одно изображение в модель, чтобы проверить, может ли мой графический процессор справиться с этим. Следовательно, уменьшение размера партии не является возможным решением.