Sync Batchnorm

[dougsouza]

Pytorch guys recently released an official SyncBatchnorm implementation. It requires a specific setup where we use torch.parallel.DistributedDataParallel(...) instead of nn.DataParallel(...) and launch a separate process for each GPU.

I wrote a small step-by-step here: https://github.com/dougsouza/pytorch-sync-batchnorm-example.

In my experiments SyncBatchnorm worked well. Also, using torch.parallel.DistributedDataParallel(...) with one process per GPU provides a huge speed up in training. The gain of adding more GPUs is almost linear, it performs a lot faster than nn.DataParallel(...). I believe you could reduce training time drastically by switching to torch.parallel.DistributedDataParallel(...).

BTW, thanks for this implementation!

[ajbrock]

If you're interested in testing and adding it to this codebase, please feel free to make a PR.