1) How does batch normalization layer work with multi_gpu_model?
For N GPUs, there are N copies of model, one on each GPU. For each copy, forward and backward passes are executed for a sub-batch (each sub-batch is 1/Nth of a batch). This means, batch normalization is actually sub-batch normalization, there is no access to the rest of batch.
# This `fit` call will be distributed on 8 GPUs.
# Since the batch size is 256, each GPU will process 32 samples.
parallel_model.fit(x, y, epochs=20, batch_size=256)
2) Which batch normalization parameters are saved when saving a model
(since when using multiple-gpus in Keras original model must be saved,
as suggested here)?
One unified (template) weight is saved instead of N different weights. Each GPU calculates a different gradient due to receiving a different sub-batch. Then, either (1) weights are updated separately on each GPU and become synchronized periodically, or (2) N outputs/gradients are aggregated on the template model (on CPU), and then new weights are broadcasted back to GPUs. In both cases, there is a unified (template) model, although in the first case, the template model may not always have the latest weights since it is kept out of synchronization and gets updated occasionally.
Some researchers have proposed a specific synchronizing technique for batch normalization to utilize the whole batch instead of a sub-batch. They state:
Standard Implementations of BN in public frameworks (suck as Caffe,
MXNet, Torch, TF, PyTorch) are unsynchronized, which means that the
data are normalized within each GPU.
Because of weights synchronization we cannot expect a linear speed-up w.r.t. number of GPUs.
Some technicalities of
multi_gpu_model are discussed in this github issue.
multi_gpu_model has a speed gain when weights are sparse (in comparison to Dense layers), otherwise weights synchronization becomes a bottleneck.
Also, here is an example of GPU-GPU weight synchronization flow from Nvidia: