HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

ultralytics 8.2.61 fix num_threads for CPU training (#14544)

Browse source 
Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Laughing 2024-07-21 02:13:45 +08:00 • committed by GitHub
parent 8094df3c47
commit 47ff2b4a76
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 239 additions and 110 deletions
Download patch file Download diff file Expand all files Collapse all files

3
ultralytics/utils/torch_utils.py
View file

@ -21,6 +21,7 @@ from ultralytics.utils import (
DEFAULT_CFG_DICT,
DEFAULT_CFG_KEYS,
LOGGER,
NUM_THREADS,
PYTHON_VERSION,
TORCHVISION_VERSION,
__version__,
@ -172,6 +173,8 @@ def select_device(device="", batch=0, newline=False, verbose=True):
s += f"CPU ({get_cpu_info()})\n"
arg = "cpu"
if arg in {"cpu", "mps"}:
torch.set_num_threads(NUM_THREADS) # reset OMP_NUM_THREADS for cpu training
if verbose:
LOGGER.info(s if newline else s.rstrip())
return torch.device(arg)