Керас: Как использовать multi-GPU (multi_gpu_models) в настройке кодера-декодера?

Question

У меня есть две модели: кодер и декодер, которые я объединяю в вариационный автоэнкодер.

Этот код работает в случае с одним графическим процессором:

enc = Model(inputs = E_input, outputs = z_coord, name='encoder')
dec = Model(inputs = D_input, outputs = D_output, name = 'decoder')

vae_outputs = dec(enc(E_input))
endec = Model(inputs = E_input, outputs = vae_outputs, name='cnn vae mlp')

Как мне лучше настроить его для работы на нескольких графических процессорах? Следующее выглядит разумно, но не работает. Это также не работает, делайте make enc и dec в multi_gpu_models (например, что закомментировано).

GPUS - это константа, равная количеству физических графических процессоров.

    with tf.device("/cpu:0"):
        enc = Model(inputs = E_input, outputs = z_coord, name='encoder')
        dec = Model(inputs = D_input, outputs = D_output, name = 'decoder')

        vae_outputs = dec(enc(E_input))
        endec = Model(inputs = E_input, outputs = vae_outputs, name='cnn vae mlp')


    # enc = multi_gpu_model(enc, gpus = GPUS)
    # dec = multi_gpu_model(dec, gpus = GPUS)
    endec = multi_gpu_model(endec, gpus = GPUS)