Fix RayTune for < 8 CPUs (#3102)

single_channel
Glenn Jocher 1 year ago committed by GitHub
parent 3ae81ee9d1
commit e88c89d269
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

@ -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,

@ -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:

Loading…
Cancel
Save