merge multi-gpu and single gpu mining mode

Author: yzbx

det-yolov5-tmi/mining/mining_cald.py

@@ -19,18 +19,35 @@
 import cv2
 import numpy as np
 import torch.utils.data as td
+from mining.data_augment import cutout, horizontal_flip, intersect, resize, rotate
+from nptyping import NDArray
 from scipy.stats import entropy
 from torch.utils.data._utils.collate import default_collate
-
-from mining.data_augment import cutout, horizontal_flip, resize, rotate
-from mining.mining_cald import get_ious
 from utils.augmentations import letterbox
+from utils.ymir_yolov5 import BBOX
 
 LOCAL_RANK = int(os.getenv('LOCAL_RANK', -1))  # https://pytorch.org/docs/stable/elastic/run.html
 RANK = int(os.getenv('RANK', -1))
 WORLD_SIZE = int(os.getenv('WORLD_SIZE', 1))
 
 
+def get_ious(boxes1: BBOX, boxes2: BBOX) -> NDArray:
+    """
+    args:
+        boxes1: np.array, (N, 4), xyxy
+        boxes2: np.array, (M, 4), xyxy
+    return:
+        iou: np.array, (N, M)
+    """
+    area1 = (boxes1[:, 2] - boxes1[:, 0]) * (boxes1[:, 3] - boxes1[:, 1])
+    area2 = (boxes2[:, 2] - boxes2[:, 0]) * (boxes2[:, 3] - boxes2[:, 1])
+    iner_area = intersect(boxes1, boxes2)
+    area1 = area1.reshape(-1, 1).repeat(area2.shape[0], axis=1)
+    area2 = area2.reshape(1, -1).repeat(area1.shape[0], axis=0)
+    iou = iner_area / (area1 + area2 - iner_area + 1e-14)
+    return iou
+
+
 def preprocess(img, img_size, stride):
     img1 = letterbox(img, img_size, stride=stride, auto=False)[0]
 

det-yolov5-tmi/mining/mining_entropy.py

@@ -174,7 +174,7 @@ def run(ymir_cfg: edict, ymir_yolov5: YmirYolov5):
         if RANK in [-1, 0]:
             ymir_yolov5.write_monitor_logger(stage=YmirStage.TASK, p=idx * batch_size_per_gpu / dataset_size)
 
-    torch.save(mining_results, f'/out/mining_results_{RANK}.pt')
+    torch.save(mining_results, f'/out/mining_results_{max(0,RANK)}.pt')
 
 
 def main() -> int:
@@ -191,9 +191,7 @@ def main() -> int:
 
     # wait all process to save the mining result
     if LOCAL_RANK != -1:
-        print(f'rank: {RANK}, sync start before merge')
         dist.barrier()
-        print(f'rank: {RANK}, sync finished before merge')
 
     if RANK in [0, -1]:
         results = []

det-yolov5-tmi/mining/util.py

@@ -155,7 +155,7 @@ def run(ymir_cfg: edict, ymir_yolov5: YmirYolov5):
             image_file = batch['image_file'][inner_idx]
             mining_results[image_file] = batch_consistency[inner_idx]
 
-    torch.save(mining_results, f'/out/mining_results_{RANK}.pt')
+    torch.save(mining_results, f'/out/mining_results_{max(0,RANK)}.pt')
 
 
 def main() -> int:

det-yolov5-tmi/mining/ymir_mining_aldd.py

@@ -14,15 +14,12 @@
 import torch.distributed as dist
 import torch.utils.data as td
 from easydict import EasyDict as edict
+from mining.util import YmirDataset, load_image_file
 from tqdm import tqdm
+from utils.ymir_yolov5 import YmirYolov5
 from ymir_exc import result_writer as rw
 from ymir_exc.util import YmirStage, get_merged_config
 
-from mining.util import (YmirDataset, collate_fn_with_fake_ann, load_image_file, load_image_file_with_ann,
-                         update_consistency)
-from utils.general import scale_coords
-from utils.ymir_yolov5 import YmirYolov5
-
 LOCAL_RANK = int(os.getenv('LOCAL_RANK', -1))  # https://pytorch.org/docs/stable/elastic/run.html
 RANK = int(os.getenv('RANK', -1))
 WORLD_SIZE = int(os.getenv('WORLD_SIZE', 1))
@@ -58,9 +55,7 @@ def run(ymir_cfg: edict, ymir_yolov5: YmirYolov5):
                                           pin_memory=ymir_yolov5.pin_memory,
                                           drop_last=False)
 
-    results = []
     mining_results = dict()
-    beta = 1.3
     dataset_size = len(images_rank)
     pbar = tqdm(origin_dataset_loader) if RANK == 0 else origin_dataset_loader
     for idx, batch in enumerate(pbar):
@@ -73,18 +68,16 @@ def run(ymir_cfg: edict, ymir_yolov5: YmirYolov5):
 
         if RANK in [-1, 0]:
             ymir_yolov5.write_monitor_logger(stage=YmirStage.TASK, p=idx * batch_size_per_gpu / dataset_size)
-        preprocess_image_shape = batch['image'].shape[2:]
         for inner_idx, det in enumerate(pred):  # per image
-            result_per_image = []
             image_file = batch['image_file'][inner_idx]
             if len(det):
                 conf = det[:, 4].data.cpu().numpy()
-                mining_results[image_file] = -np.sum(conf*np.log2(conf))
+                mining_results[image_file] = -np.sum(conf * np.log2(conf))
             else:
                 mining_results[image_file] = -10
                 continue
 
-    torch.save(mining_results, f'/out/mining_results_{RANK}.pt')
+    torch.save(mining_results, f'/out/mining_results_{max(0,RANK)}.pt')
 
 
 def main() -> int:
@@ -99,7 +92,7 @@ def main() -> int:
     run(ymir_cfg, ymir_yolov5)
 
     # wait all process to save the mining result
-    if LOCAL_RANK != -1:
+    if WORLD_SIZE > 1:
         dist.barrier()
 
     if RANK in [0, -1]:
@@ -112,10 +105,6 @@ def main() -> int:
             for img_file, score in result.items():
                 ymir_mining_result.append((img_file, score))
         rw.write_mining_result(mining_result=ymir_mining_result)
-
-    if LOCAL_RANK != -1:
-        print(f'rank: {RANK}, start destroy process group')
-        # dist.destroy_process_group()
     return 0