Implement syncbn for TensorFlow

[edwardyehuang]

Add SyncBN implementation (via synchronized argument in layers.BatchNormalization) for TensorFlow. Note that, I don't know how to write the test for it in Keras online test (it requires multi-gpus).

#18667

[edwardyehuang]

Also, I am not sure if it is better to move the framework-specified code (e.g. TensorFlow) to backbend, the current distribution_lib.py states "!!!DO NOT USE!!!"

[codecov-commenter]

Codecov Report

Attention: 27 lines in your changes are missing coverage. Please review.

Comparison is base (cb65582) 78.53% compared to head (e7a442f) 78.47%.

Additional details and impacted files

@@            Coverage Diff             @@
##           master   #18671      +/-   ##
==========================================
- Coverage   78.53%   78.47%   -0.07%     
==========================================
  Files         335      335              
  Lines       32943    32975      +32     
  Branches     6450     6454       +4     
==========================================
+ Hits        25873    25878       +5     
- Misses       5512     5538      +26     
- Partials     1558     1559       +1     

Flag	Coverage Δ
keras	78.37% <18.18%> (-0.07%)	⬇️
keras-jax	63.39% <18.18%> (-0.05%)	⬇️
keras-numpy	57.70% <18.18%> (-0.05%)	⬇️
keras-tensorflow	64.49% <18.18%> (-0.05%)	⬇️
keras-torch	65.20% <18.18%> (-0.05%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files	Coverage Δ
keras/layers/normalization/batch_normalization.py	74.76% <18.18%> (-25.24%)	⬇️

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[fchollet]

Thanks for the PR!

@qlzh727 can you advise on how to proceed for testing the feature with 2 devices?

[fchollet]

LGTM, thanks. Will post-process this on our side.

[qlzh727]

For sync batch norm test, user can config 2 virtual cpus and use them for mirrored strategy test.

https://www.tensorflow.org/api_docs/python/tf/config/set_logical_device_configuration

[fchollet]

@qlzh727 the test for this feature is now moved to nn_test.py:test_moments_sync: https://github.com/keras-team/keras/blob/master/keras/ops/nn_test.py#L1441

How should we modify the test to ensure correctness? (right now I think the sync branch is actually never run)

[qlzh727]

ack, I will add a test for that. The sync logic will only have a difference in the distribution setting.