chore(core): update launch backend to subprocess (#158)

tools/train.py

@@ -12,7 +12,7 @@ import torch.backends.cudnn as cudnn
 
 from yolox.core import Trainer, launch
 from yolox.exp import get_exp
-from yolox.utils import configure_nccl
+from yolox.utils import configure_nccl, get_num_devices
 
 
 def make_parser():
@@ -106,8 +106,8 @@ if __name__ == "__main__":
     exp = get_exp(args.exp_file, args.name)
     exp.merge(args.opts)
 
-    num_gpu = torch.cuda.device_count() if args.devices is None else args.devices
-    assert num_gpu <= torch.cuda.device_count()
+    num_gpu = get_num_devices() if args.devices is None else args.devices
+    assert num_gpu <= get_num_devices()
 
     dist_url = "auto" if args.dist_url is None else args.dist_url
     launch(

yolox/core/launch.py

@@ -55,15 +55,17 @@ def launch(
             port = _find_free_port()
             dist_url = f"tcp://127.0.0.1:{port}"
 
-        mp.spawn(
-            _distributed_worker,
-            nprocs=num_gpus_per_machine,
-            args=(
-                main_func, world_size, num_gpus_per_machine,
-                machine_rank, backend, dist_url, args
-            ),
-            daemon=False,
-        )
+        processes = []
+        for rank in range(num_gpus_per_machine):
+            p = mp.Process(
+                target=_distributed_worker, 
+                args=(
+                    rank, main_func, world_size, num_gpus_per_machine,
+                    machine_rank, backend, dist_url, args))
+            p.start()
+            processes.append(p)
+        for p in processes:
+            p.join()
     else:
         main_func(*args)
 

yolox/data/data_augment.py

@@ -4,9 +4,6 @@
 """
 Data augmentation functionality. Passed as callable transformations to
 Dataset classes.
-
-The data augmentation procedures were interpreted from @weiliu89's SSD paper
-http://arxiv.org/abs/1512.02325
 """
 
 import math
@@ -17,6 +14,8 @@ import numpy as np
 
 import torch
 
+from yolox.utils import xyxy2cxcywh
+
 
 def augment_hsv(img, hgain=0.015, sgain=0.7, vgain=0.4):
     r = np.random.uniform(-1, 1, 3) * [hgain, sgain, vgain] + 1  # random gains
@@ -197,20 +196,11 @@ class TrainTransform:
     def __call__(self, image, targets, input_dim):
         boxes = targets[:, :4].copy()
         labels = targets[:, 4].copy()
-        if targets.shape[1] > 5:
-            mixup = True
-            ratios = targets[:, -1].copy()
-            ratios_o = targets[:, -1].copy()
-        else:
-            mixup = False
-            ratios = None
-            ratios_o = None
-        lshape = 6 if mixup else 5
         if len(boxes) == 0:
-            targets = np.zeros((self.max_labels, lshape), dtype=np.float32)
+            targets = np.zeros((self.max_labels, 5), dtype=np.float32)
             image, r_o = preproc(image, input_dim, self.means, self.std)
             image = np.ascontiguousarray(image, dtype=np.float32)
-            return image, targets
+            return torch.as_tensor(image), torch.as_tensor(targets)
 
         image_o = image.copy()
         targets_o = targets.copy()
@@ -218,58 +208,36 @@ class TrainTransform:
         boxes_o = targets_o[:, :4]
         labels_o = targets_o[:, 4]
         # bbox_o: [xyxy] to [c_x,c_y,w,h]
-        b_x_o = (boxes_o[:, 2] + boxes_o[:, 0]) * 0.5
-        b_y_o = (boxes_o[:, 3] + boxes_o[:, 1]) * 0.5
-        b_w_o = (boxes_o[:, 2] - boxes_o[:, 0]) * 1.0
-        b_h_o = (boxes_o[:, 3] - boxes_o[:, 1]) * 1.0
-        boxes_o[:, 0] = b_x_o
-        boxes_o[:, 1] = b_y_o
-        boxes_o[:, 2] = b_w_o
-        boxes_o[:, 3] = b_h_o
+        boxes_o = xyxy2cxcywh(boxes_o)
 
         image_t = _distort(image)
         image_t, boxes = _mirror(image_t, boxes)
         height, width, _ = image_t.shape
         image_t, r_ = preproc(image_t, input_dim, self.means, self.std)
-        boxes = boxes.copy()
+        boxes = xyxy2cxcywh(boxes)
         # boxes [xyxy] 2 [cx,cy,w,h]
-        b_x = (boxes[:, 2] + boxes[:, 0]) * 0.5
-        b_y = (boxes[:, 3] + boxes[:, 1]) * 0.5
-        b_w = (boxes[:, 2] - boxes[:, 0]) * 1.0
-        b_h = (boxes[:, 3] - boxes[:, 1]) * 1.0
-        boxes[:, 0] = b_x
-        boxes[:, 1] = b_y
-        boxes[:, 2] = b_w
-        boxes[:, 3] = b_h
 
         boxes *= r_
 
         mask_b = np.minimum(boxes[:, 2], boxes[:, 3]) > 8
         boxes_t = boxes[mask_b]
-        labels_t = labels[mask_b].copy()
-        if mixup:
-            ratios_t = ratios[mask_b].copy()
+        labels_t = labels[mask_b]
 
         if len(boxes_t) == 0:
             image_t, r_o = preproc(image_o, input_dim, self.means, self.std)
             boxes_o *= r_o
             boxes_t = boxes_o
             labels_t = labels_o
-            ratios_t = ratios_o
 
         labels_t = np.expand_dims(labels_t, 1)
-        if mixup:
-            ratios_t = np.expand_dims(ratios_t, 1)
-            targets_t = np.hstack((labels_t, boxes_t, ratios_t))
-        else:
-            targets_t = np.hstack((labels_t, boxes_t))
-        padded_labels = np.zeros((self.max_labels, lshape))
+        targets_t = np.hstack((labels_t, boxes_t))
+        padded_labels = np.zeros((self.max_labels, 5))
         padded_labels[range(len(targets_t))[: self.max_labels]] = targets_t[
             : self.max_labels
         ]
         padded_labels = np.ascontiguousarray(padded_labels, dtype=np.float32)
         image_t = np.ascontiguousarray(image_t, dtype=np.float32)
-        return image_t, padded_labels
+        return torch.as_tensor(image_t), torch.as_tensor(padded_labels)
 
 
 class ValTransform:
@@ -298,4 +266,4 @@ class ValTransform:
     # assume input is cv2 img for now
     def __call__(self, img, res, input_size):
         img, _ = preproc(img, input_size, self.means, self.std, self.swap)
-        return torch.from_numpy(img), torch.zeros(1, 5)
+        return torch.tensor_as(img), torch.zeros(1, 5)

yolox/data/datasets/mosaicdetection.py

@@ -220,4 +220,4 @@ class MosaicDetection(Dataset):
             origin_img = origin_img.astype(np.float32)
             origin_img = 0.5 * origin_img + 0.5 * padded_cropped_img.astype(np.float32)
 
-        return origin_img.astype(np.uint8), origin_labels
+        return origin_img, origin_labels

yolox/utils/boxes.py

@@ -9,7 +9,7 @@ import torchvision
 
 __all__ = [
     "filter_box", "postprocess", "bboxes_iou", "matrix_iou",
-    "adjust_box_anns", "xyxy2xywh",
+    "adjust_box_anns", "xyxy2xywh", "xyxy2cxcywh",
 ]
 
 
@@ -113,3 +113,11 @@ def xyxy2xywh(bboxes):
     bboxes[:, 2] = bboxes[:, 2] - bboxes[:, 0]
     bboxes[:, 3] = bboxes[:, 3] - bboxes[:, 1]
     return bboxes
+
+
+def xyxy2cxcywh(bboxes):
+    bboxes[:, 2] = bboxes[:, 2] - bboxes[:, 0]
+    bboxes[:, 3] = bboxes[:, 3] - bboxes[:, 1]
+    bboxes[:, 0] = bboxes[:, 0] + bboxes[:, 2] * 0.5
+    bboxes[:, 1] = bboxes[:, 1] + bboxes[:, 3] * 0.5
+    return bboxes

yolox/utils/metric.py

@@ -13,12 +13,19 @@ import torch
 __all__ = [
     "AverageMeter",
     "MeterBuffer",
+    "get_num_devices",
     "get_total_and_free_memory_in_Mb",
     "occupy_mem",
     "gpu_mem_usage",
 ]
 
 
+def get_num_devices():
+    devices_list_info = os.popen("nvidia-smi -L")
+    devices_list_info = devices_list_info.read().strip().split("\n")
+    return len(devices_list_info)
+
+
 def get_total_and_free_memory_in_Mb(cuda_device):
     devices_info_str = os.popen(
         "nvidia-smi --query-gpu=memory.total,memory.used --format=csv,nounits,noheader"

yolox/utils/setup_env.py

@@ -48,3 +48,5 @@ def configure_module(ulimit_value=8192):
     except Exception:
         # cv2 version mismatch might rasie exceptions.
         pass
+
+    os.environ["OMP_NUM_THREADS"] = str(1)