1) How does the batch normalization layer work with multi_gpu_model?

Is it calculated separately on each GPU, or is somehow synchronized between GPUs?

2) Which batch normalization parameters are saved when saving a model? (Since when using multiple-gpus in Keras, the original model must be saved, as suggested here)?

In the docs of multi_gpu_model is says:

Specifically, this function implements single-machine multi-GPU data parallelism.

What does it mean for Batch Normalization?


1 Answer 1


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.

Extra remarks

  1. 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.

  2. Because of weights synchronization we cannot expect a linear speed-up w.r.t. number of GPUs.

  3. 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.

  4. Also, here is an example of GPU-GPU weight synchronization flow from Nvidia:

  • $\begingroup$ How do you know that Keras does weights synchronization since in the docs of multi_gpu_model function it only says that original model and the parallel model share weights (# Save model via the template model (which shares the same weights)) ? $\endgroup$ Mar 25, 2019 at 9:49
  • $\begingroup$ @AntonioJurić you are right, that's an inaccurate statement, because if all N copies have the same weight at all times, what the .fit is doing on separate GPUs and separate data? So weights are different temporarily and become unified periodically. $\endgroup$
    – Esmailian
    Mar 25, 2019 at 11:15
  • $\begingroup$ @AntonioJurić I was wrong in my previous comment. They could have the same weights at all times. Updated. $\endgroup$
    – Esmailian
    Mar 29, 2019 at 15:35

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

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