Fix bug: optimize bn when distributed is True

tianweiy · Apr 1, 2021 · 607f01a · 607f01a
1 parent 0384583
commit 607f01a
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/det3d/torchie/apis/train.py b/det3d/torchie/apis/train.py
@@ -265,7 +265,8 @@ def train_detector(model, dataset, cfg, distributed=False, validate=False, logge
 
     total_steps = cfg.total_epochs * len(data_loaders[0])
     # print(f"total_steps: {total_steps}")
-
+    if distributed:
+        model = apex.parallel.convert_syncbn_model(model)
     if cfg.lr_config.type == "one_cycle":
         # build trainer
         optimizer = build_one_cycle_optimizer(model, cfg.optimizer)
@@ -281,7 +282,6 @@ def train_detector(model, dataset, cfg, distributed=False, validate=False, logge
 
     # put model on gpus
     if distributed:
-        model = apex.parallel.convert_syncbn_model(model)
         model = DistributedDataParallel(
             model.cuda(cfg.local_rank),
             device_ids=[cfg.local_rank],