fix 30min dist.barrier() time-out errors

Author: yzbx

det-mmdetection-tmi/ymir_mining.py

@@ -283,6 +283,10 @@ def mining(self):
         beta = 1.3
         mining_result = []
         for asset_path in tbar:
+            # batch-level sync, avoid 30min time-out error
+            if LOCAL_RANK != -1:
+                dist.barrier()
+
             img = cv2.imread(asset_path)
             # xyxy,conf,cls
             result = self.predict(img)

det-yolov5-tmi/mining/ymir_infer.py

@@ -58,6 +58,10 @@ def run(ymir_cfg: edict, ymir_yolov5: YmirYolov5):
     monitor_gap = max(1, dataset_size // 1000 // batch_size_per_gpu)
     pbar = tqdm(origin_dataset_loader) if RANK == 0 else origin_dataset_loader
     for idx, batch in enumerate(pbar):
+        # batch-level sync, avoid 30min time-out error
+        if LOCAL_RANK != -1:
+            dist.barrier()
+
         with torch.no_grad():
             pred = ymir_yolov5.forward(batch['image'].float().to(device), nms=True)
 

det-yolov5-tmi/mining/ymir_mining_aldd.py

@@ -157,6 +157,10 @@ def run(ymir_cfg: edict, ymir_yolov5: YmirYolov5):
     pbar = tqdm(origin_dataset_loader) if RANK == 0 else origin_dataset_loader
     miner = ALDD(ymir_cfg)
     for idx, batch in enumerate(pbar):
+        # batch-level sync, avoid 30min time-out error
+        if LOCAL_RANK != -1:
+            dist.barrier()
+
         with torch.no_grad():
             featuremap_output = ymir_yolov5.model.model(batch['image'].float().to(device))[1]
             unc_scores = miner.compute_aldd_score(featuremap_output, ymir_yolov5.img_size)
@@ -200,7 +204,7 @@ def main() -> int:
 
     if LOCAL_RANK != -1:
         print(f'rank: {RANK}, start destroy process group')
-    dist.destroy_process_group()
+        dist.destroy_process_group()
     return 0
 
 

det-yolov5-tmi/mining/ymir_mining_cald.py

@@ -62,6 +62,10 @@ def run(ymir_cfg: edict, ymir_yolov5: YmirYolov5):
     dataset_size = len(images_rank)
     pbar = tqdm(origin_dataset_loader) if RANK == 0 else origin_dataset_loader
     for idx, batch in enumerate(pbar):
+        # batch-level sync, avoid 30min time-out error
+        if LOCAL_RANK != -1:
+            dist.barrier()
+
         with torch.no_grad():
             pred = ymir_yolov5.forward(batch['image'].float().to(device), nms=True)
 
@@ -98,6 +102,10 @@ def run(ymir_cfg: edict, ymir_yolov5: YmirYolov5):
     monitor_gap = max(1, dataset_size // 1000 // batch_size_per_gpu)
     pbar = tqdm(aug_dataset_loader) if RANK == 0 else aug_dataset_loader
     for idx, batch in enumerate(pbar):
+        # batch-level sync, avoid 30min time-out error
+        if LOCAL_RANK != -1:
+            dist.barrier()
+
         if idx % monitor_gap == 0 and RANK in [-1, 0]:
             ymir_yolov5.write_monitor_logger(stage=YmirStage.TASK, p=idx * batch_size_per_gpu / dataset_size)