merge single-gpu and mutiple gpu infer

Author: yzbx

det-yolov5-tmi/mypy.ini

@@ -2,7 +2,6 @@
 ignore_missing_imports = True
 disallow_untyped_defs = False
 exclude = [utils/general.py, models/*.py, utils/*.py]
-files = mining/*.py, utils/ymir_yolov5.py, start.py, train.py
 
 [mypy-torch.*]
 ignore_errors = True

det-yolov5-tmi/ymir/README.md

@@ -9,9 +9,9 @@ docker build -t ymir/ymir-executor:ymir1.1.0-cuda111-yolov5-tmi --build-arg SERV
 
 ## main change log
 
-- add `start.py` and `utils/ymir_yolov5.py` for train/infer/mining
+- add `start.py` and `ymir/ymir_yolov5.py` for train/infer/mining
 
-- add `utils/ymir_yolov5.py` for useful functions
+- add `ymir/ymir_yolov5.py` for useful functions
 
     - `get_merged_config()` add ymir path config `cfg.yaml` and hyper-parameter `cfg.param`
 

det-yolov5-tmi/ymir/mining/data_augment.py

@@ -8,8 +8,7 @@
 import cv2
 import numpy as np
 from nptyping import NDArray
-
-from utils.ymir_yolov5 import BBOX, CV_IMAGE
+from ymir.ymir_yolov5 import BBOX, CV_IMAGE
 
 
 def intersect(boxes1: BBOX, boxes2: BBOX) -> NDArray:

det-yolov5-tmi/ymir/mining/util.py

@@ -19,12 +19,12 @@
 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 utils.augmentations import letterbox
-from utils.ymir_yolov5 import BBOX
+from ymir.mining.data_augment import cutout, horizontal_flip, intersect, resize, rotate
+from ymir.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))

det-yolov5-tmi/ymir/mining/ymir_infer.py

@@ -12,10 +12,10 @@
 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.general import scale_coords
-from utils.ymir_yolov5 import YmirYolov5
+from ymir.mining.util import YmirDataset, load_image_file
+from ymir.ymir_yolov5 import YmirYolov5
 from ymir_exc import result_writer as rw
 from ymir_exc.util import YmirStage, get_merged_config
 
@@ -44,7 +44,10 @@ def run(ymir_cfg: edict, ymir_yolov5: YmirYolov5):
 
     max_barrier_times = len(images) // max(1, WORLD_SIZE) // batch_size_per_gpu
     # origin dataset
-    images_rank = images[RANK::WORLD_SIZE]
+    if RANK != -1:
+        images_rank = images[RANK::WORLD_SIZE]
+    else:
+        images_rank = images
     origin_dataset = YmirDataset(images_rank, load_fn=load_fn)
     origin_dataset_loader = td.DataLoader(origin_dataset,
                                           batch_size=batch_size_per_gpu,
@@ -80,7 +83,7 @@ def run(ymir_cfg: edict, ymir_yolov5: YmirYolov5):
                 result_per_image.append(det)
             results.append(dict(image_file=image_file, result=result_per_image))
 
-    torch.save(results, f'/out/infer_results_{RANK}.pt')
+    torch.save(results, f'/out/infer_results_{max(0,RANK)}.pt')
 
 
 def main() -> int:

det-yolov5-tmi/ymir/mining/ymir_mining_aldd.py

@@ -17,9 +17,9 @@
 import torch.nn.functional as F
 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.mining.util import YmirDataset, load_image_file
+from ymir.ymir_yolov5 import YmirYolov5
 from ymir_exc import result_writer as rw
 from ymir_exc.util import YmirStage, get_merged_config
 

det-yolov5-tmi/ymir/mining/ymir_mining_cald.py

@@ -14,11 +14,11 @@
 import torch.distributed as dist
 import torch.utils.data as td
 from easydict import EasyDict as edict
-from mining.util import (YmirDataset, collate_fn_with_fake_ann, load_image_file, load_image_file_with_ann,
-                         update_consistency)
 from tqdm import tqdm
 from utils.general import scale_coords
-from utils.ymir_yolov5 import YmirYolov5
+from ymir.mining.util import (YmirDataset, collate_fn_with_fake_ann, load_image_file, load_image_file_with_ann,
+                              update_consistency)
+from ymir.ymir_yolov5 import YmirYolov5
 from ymir_exc import result_writer as rw
 from ymir_exc.util import YmirStage, get_merged_config
 
@@ -47,7 +47,10 @@ def run(ymir_cfg: edict, ymir_yolov5: YmirYolov5):
 
     max_barrier_times = (len(images) // max(1, WORLD_SIZE)) // batch_size_per_gpu
     # origin dataset
-    images_rank = images[RANK::WORLD_SIZE]
+    if RANK != -1:
+        images_rank = images[RANK::WORLD_SIZE]
+    else:
+        images_rank = images
     origin_dataset = YmirDataset(images_rank, load_fn=load_fn)
     origin_dataset_loader = td.DataLoader(origin_dataset,
                                           batch_size=batch_size_per_gpu,

det-yolov5-tmi/ymir/mining/ymir_mining_entropy.py

@@ -14,9 +14,9 @@
 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.mining.util import YmirDataset, load_image_file
+from ymir.ymir_yolov5 import YmirYolov5
 from ymir_exc import result_writer as rw
 from ymir_exc.util import YmirStage, get_merged_config
 
@@ -45,7 +45,10 @@ def run(ymir_cfg: edict, ymir_yolov5: YmirYolov5):
 
     max_barrier_times = (len(images) // max(1, WORLD_SIZE)) // batch_size_per_gpu
     # origin dataset
-    images_rank = images[RANK::WORLD_SIZE]
+    if RANK != -1:
+        images_rank = images[RANK::WORLD_SIZE]
+    else:
+        images_rank = images
     origin_dataset = YmirDataset(images_rank, load_fn=load_fn)
     origin_dataset_loader = td.DataLoader(origin_dataset,
                                           batch_size=batch_size_per_gpu,

det-yolov5-tmi/ymir/mining/ymir_mining_random.py

@@ -13,7 +13,7 @@
 import torch.distributed as dist
 from easydict import EasyDict as edict
 from tqdm import tqdm
-from utils.ymir_yolov5 import YmirYolov5
+from ymir.ymir_yolov5 import YmirYolov5
 from ymir_exc import result_writer as rw
 from ymir_exc.util import YmirStage, get_merged_config
 
@@ -31,7 +31,10 @@ def run(ymir_cfg: edict, ymir_yolov5: YmirYolov5):
     with open(ymir_cfg.ymir.input.candidate_index_file, 'r') as f:
         images = [line.strip() for line in f.readlines()]
 
-    images_rank = images[RANK::WORLD_SIZE]
+    if RANK != -1:
+        images_rank = images[RANK::WORLD_SIZE]
+    else:
+        images_rank = images
     mining_results = dict()
     dataset_size = len(images_rank)
     pbar = tqdm(images_rank) if RANK == 0 else images_rank

det-yolov5-tmi/ymir/start.py

@@ -6,7 +6,7 @@
 import cv2
 from easydict import EasyDict as edict
 from models.experimental import attempt_download
-from utils.ymir_yolov5 import YmirYolov5, convert_ymir_to_yolov5, get_weight_file
+from ymir.ymir_yolov5 import YmirYolov5, convert_ymir_to_yolov5, get_weight_file
 from ymir_exc import dataset_reader as dr
 from ymir_exc import env, monitor
 from ymir_exc import result_writer as rw
@@ -124,10 +124,10 @@ def _run_mining(cfg: edict, task_idx: int = 0, task_num: int = 1) -> None:
         raise Exception(f'unknown mining algorithm {mining_algorithm}, not in {support_mining_algorithms}')
 
     if gpu_count <= 1:
-        command = f'python3 mining/ymir_mining_{mining_algorithm}.py'
+        command = f'python3 ymir/mining/ymir_mining_{mining_algorithm}.py'
     else:
         port = find_free_port()
-        command = f'python3 -m torch.distributed.launch --nproc_per_node {gpu_count} --master_port {port} mining/ymir_mining_{mining_algorithm}.py'  # noqa
+        command = f'python3 -m torch.distributed.launch --nproc_per_node {gpu_count} --master_port {port} ymir/mining/ymir_mining_{mining_algorithm}.py'  # noqa
 
     logging.info(f'mining: {command}')
     subprocess.run(command.split(), check=True)
@@ -147,29 +147,13 @@ def _run_infer(cfg: edict, task_idx: int = 0, task_num: int = 1) -> None:
     gpu_count: int = len(gpu_id.split(',')) if gpu_id else 0
 
     if gpu_count <= 1:
-        N = dr.items_count(env.DatasetType.CANDIDATE)
-        infer_result = dict()
-        model = YmirYolov5(cfg)
-        idx = -1
-
-        monitor_gap = max(1, N // 100)
-        for asset_path, _ in dr.item_paths(dataset_type=env.DatasetType.CANDIDATE):
-            img = cv2.imread(asset_path)
-            result = model.infer(img)
-            infer_result[asset_path] = result
-            idx += 1
-
-            if idx % monitor_gap == 0:
-                percent = get_ymir_process(stage=YmirStage.TASK, p=idx / N, task_idx=task_idx, task_num=task_num)
-                monitor.write_monitor_logger(percent=percent)
-
-        rw.write_infer_result(infer_result=infer_result)
+        command = 'python3 ymir/mining/ymir_infer.py'
     else:
         port = find_free_port()
-        command = f'python3 -m torch.distributed.launch --nproc_per_node {gpu_count} --master_port {port} mining/ymir_infer.py'  # noqa
+        command = f'python3 -m torch.distributed.launch --nproc_per_node {gpu_count} --master_port {port} ymir/mining/ymir_infer.py'  # noqa
 
-        logging.info(f'infer: {command}')
-        subprocess.run(command.split(), check=True)
+    logging.info(f'infer: {command}')
+    subprocess.run(command.split(), check=True)
 
     monitor.write_monitor_logger(
         percent=get_ymir_process(stage=YmirStage.POSTPROCESS, p=1.0, task_idx=task_idx, task_num=task_num))