From dda78692050e01d0c546c790be4d867b86e937b4 Mon Sep 17 00:00:00 2001 From: Andreas Moldskred Date: Thu, 2 May 2024 18:21:18 -0700 Subject: [PATCH] BUGFIX: Dataloading CPU bottleneck (#11195) Co-authored-by: Glenn Jocher Co-authored-by: UltralyticsAssistant --- ultralytics/data/build.py | 3 ++- ultralytics/utils/__init__.py | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/ultralytics/data/build.py b/ultralytics/data/build.py index ce8785d680..d0773737a2 100644 --- a/ultralytics/data/build.py +++ b/ultralytics/data/build.py @@ -20,7 +20,7 @@ from ultralytics.data.loaders import ( autocast_list, ) from ultralytics.data.utils import IMG_FORMATS, VID_FORMATS -from ultralytics.utils import RANK, colorstr +from ultralytics.utils import NUM_THREADS, RANK, colorstr from ultralytics.utils.checks import check_file from .dataset import GroundingDataset, YOLODataset, YOLOMultiModalDataset @@ -81,6 +81,7 @@ def seed_worker(worker_id): # noqa worker_seed = torch.initial_seed() % 2**32 np.random.seed(worker_seed) random.seed(worker_seed) + os.sched_setaffinity(0, range(NUM_THREADS)) # fix https://github.com/ultralytics/ultralytics/pull/11195 def build_yolo_dataset(cfg, img_path, batch, data, mode="train", rect=False, stride=32, multi_modal=False): diff --git a/ultralytics/utils/__init__.py b/ultralytics/utils/__init__.py index 41baedca93..a82b1933a4 100644 --- a/ultralytics/utils/__init__.py +++ b/ultralytics/utils/__init__.py @@ -36,7 +36,7 @@ FILE = Path(__file__).resolve() ROOT = FILE.parents[1] # YOLO ASSETS = ROOT / "assets" # default images DEFAULT_CFG_PATH = ROOT / "cfg/default.yaml" -NUM_THREADS = min(8, max(1, os.cpu_count() - 1)) # number of YOLOv5 multiprocessing threads +NUM_THREADS = min(8, max(1, os.cpu_count() - 1)) # number of YOLO multiprocessing threads AUTOINSTALL = str(os.getenv("YOLO_AUTOINSTALL", True)).lower() == "true" # global auto-install mode VERBOSE = str(os.getenv("YOLO_VERBOSE", True)).lower() == "true" # global verbose mode TQDM_BAR_FORMAT = "{l_bar}{bar:10}{r_bar}" if VERBOSE else None # tqdm bar format