I’m using DDP(one process per GPU) to training a 3D UNet. I transfered all batchnorm layer inside network to syncbatchnorm with nn.SyncBatchNorm.convert_sync_batchnorm.
When doing validation at the end of every training epoch on rank 0, it always freeze at same validation steps. I think it is because of the syncbatchnorm layer. What is the correct way to do validation when DDP model has syncbatchnorm layer? Should I do validation on all ranks?
for epoch in range(epochs):
for step, (data, target) in enumerate(train_loader):
# ...training codes
if dist.get_rank() == 0:
# ...validation codes
Yes, you probably need to do validation on all ranks since SyncBatchNorm has collectives which are expected to run on all ranks. The validation is probably getting stuck since SyncBatchNorm on rank 0 is waiting for collectives from other ranks.
Another option is to convert the SyncBatchNorm layer to a regular BatchNorm layer and then do the validation on a single rank.
Thanks I will try it.
Actually I have another question about v1.1.0 DDP.
I tried to inference the model with syncbatchnorm layer ( Actually, it becomes batchnorm layer after load from checkpoint ). The results turned to be different between:
Only turn on evaluate mode.
Manually set track_running_stats of each BN layer to False after model.eval().
# do inference..
It is strange that the second one is much better than first one on early epochs. Is this also a version problem?
P.S. However, after more epochs training, results of two inference method are similiar but still has small differences.
Below is sample code of set_BN_track_running_stats():
def set_BN_track_running_stats(module, ifTrack=True):
if isinstance(module, torch.nn.modules.batchnorm._BatchNorm):
module.track_running_stats = ifTrack
for child in module.children():
set_BN_track_running_stats(module = child, ifTrack=ifTrack)