From e88c89d269a634b79eec946f343858fd2a4ddcf1 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Thu, 8 Jun 2023 18:50:05 +0200
Subject: [PATCH] Fix RayTune for < 8 CPUs (#3102)

---
 ultralytics/yolo/engine/model.py   | 12 +++++++-----
 ultralytics/yolo/utils/__init__.py |  2 +-
 2 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/ultralytics/yolo/engine/model.py b/ultralytics/yolo/engine/model.py
index 3529150..de3dc8d 100644
--- a/ultralytics/yolo/engine/model.py
+++ b/ultralytics/yolo/engine/model.py
@@ -9,8 +9,8 @@ from ultralytics.nn.tasks import (ClassificationModel, DetectionModel, PoseModel
                                   attempt_load_one_weight, guess_model_task, nn, yaml_model_load)
 from ultralytics.yolo.cfg import get_cfg
 from ultralytics.yolo.engine.exporter import Exporter
-from ultralytics.yolo.utils import (DEFAULT_CFG, DEFAULT_CFG_DICT, DEFAULT_CFG_KEYS, LOGGER, RANK, ROOT, callbacks,
-                                    is_git_dir, yaml_load)
+from ultralytics.yolo.utils import (DEFAULT_CFG, DEFAULT_CFG_DICT, DEFAULT_CFG_KEYS, LOGGER, NUM_THREADS, RANK, ROOT,
+                                    callbacks, is_git_dir, yaml_load)
 from ultralytics.yolo.utils.checks import check_file, check_imgsz, check_pip_update_available, check_yaml
 from ultralytics.yolo.utils.downloads import GITHUB_ASSET_STEMS
 from ultralytics.yolo.utils.torch_utils import smart_inference_mode
@@ -391,7 +391,7 @@ class YOLO:
              grace_period: int = 10,
              gpu_per_trial: int = None,
              max_samples: int = 10,
-             train_args: dict = {}):
+             train_args: dict = None):
         """
         Runs hyperparameter tuning using Ray Tune.
 
@@ -409,6 +409,8 @@ class YOLO:
         Raises:
             ModuleNotFoundError: If Ray Tune is not installed.
         """
+        if train_args is None:
+            train_args = {}
 
         try:
             from ultralytics.yolo.utils.tuner import (ASHAScheduler, RunConfig, WandbLoggerCallback, default_space,
@@ -443,7 +445,7 @@ class YOLO:
         space['data'] = data
 
         # Define the trainable function with allocated resources
-        trainable_with_resources = tune.with_resources(_tune, {'cpu': 8, 'gpu': gpu_per_trial if gpu_per_trial else 0})
+        trainable_with_resources = tune.with_resources(_tune, {'cpu': NUM_THREADS, 'gpu': gpu_per_trial or 0})
 
         # Define the ASHA scheduler for hyperparameter search
         asha_scheduler = ASHAScheduler(time_attr='epoch',
@@ -454,7 +456,7 @@ class YOLO:
                                        reduction_factor=3)
 
         # Define the callbacks for the hyperparameter search
-        tuner_callbacks = [WandbLoggerCallback(project='yolov8_tune')] if wandb else []
+        tuner_callbacks = [WandbLoggerCallback(project='YOLOv8-tune')] if wandb else []
 
         # Create the Ray Tune hyperparameter search tuner
         tuner = tune.Tuner(trainable_with_resources,
diff --git a/ultralytics/yolo/utils/__init__.py b/ultralytics/yolo/utils/__init__.py
index 77f7e20..aa08953 100644
--- a/ultralytics/yolo/utils/__init__.py
+++ b/ultralytics/yolo/utils/__init__.py
@@ -37,7 +37,7 @@ AUTOINSTALL = str(os.getenv('YOLO_AUTOINSTALL', True)).lower() == 'true'  # glob
 VERBOSE = str(os.getenv('YOLO_VERBOSE', True)).lower() == 'true'  # global verbose mode
 TQDM_BAR_FORMAT = '{l_bar}{bar:10}{r_bar}'  # tqdm bar format
 LOGGING_NAME = 'ultralytics'
-MACOS, LINUX, WINDOWS = (platform.system() == x for x in ['Darwin', 'Linux', 'Windows'])  # environment booleans
+MACOS, LINUX, WINDOWS = (platform.system() == x for x in {'Darwin', 'Linux', 'Windows'})  # environment booleans
 HELP_MSG = \
     """
     Usage examples for running YOLOv8: