TPU Batch Norm which aggregated across shards

neuroailab / tfutils

Utilities for working with tensorflow

MIT License

25 stars 8 forks source link

Closed anayebi closed 5 years ago

anayebi commented 5 years ago

Generalized cross gpu batch norm to be cross device batch norm
Supports aggregation of batch norm statistics across multiple gpus and TPU shards
Added gradient norm clipping as option to optimizer class
Unit tests pass