How do you train GANs using multiple GPUs with Keras?

Question

There are three main challenges: a) how do you save and load the optimizer state, b) how do you use multiple GPU with nested models, see below, and c), how you create a workflow to optimize GPU and CPU utilization?

Context

We have three components:

1. the discriminator
2. the generator, and
3. the GAN which has both the discriminator and the generator.

Optimizer State

Since the discriminators are included in the GAN and they also need to be used separately during training - how do you save and load GANs? Now, I save the generators and discriminators separately and recompile the GAN for each training episode, but I lose the optimizer state this way.

Multiple GPUs

This is what the API looks like:

from keras.utils import multi_gpu_model
parallel_model = multi_gpu_model(model, gpus=8)

The challenge here is the same as with optimizers. Since the discriminator is included in GANs, you can't apply the multi_gpu_model to both the discriminator and the GAN. You can add a multi_gpu_model to both the discriminator and generator before you create the GAN, but from my experience it does not scale well and leads to poor GPU utilization.

GPU and CPU utilization

The data can be preprocessed and queued using multiprocessing. Since the multi_gpu_model API does not support GANs, you need to frequently merge the weights and hop between CPUs and GPUs. Thus, I haven't found a clean way to utilize GPUs and CPUs.

Answer 1

The multi_gpu_model can be used in each of the function for generator, discriminator and gan

def create_generator():
  #network architecture 
  generator = Model(inputs=input, outputs=output)
  generator = multi_gpu_model(generator, gpus=2)
  generator.compile()
  return generator

The same can be done for discriminator and gan.