diff --git a/utils/datasets.py b/utils/datasets.py index 5a3b2110b2e0..79b871c9294b 100755 --- a/utils/datasets.py +++ b/utils/datasets.py @@ -26,8 +26,8 @@ from tqdm import tqdm from utils.augmentations import Albumentations, augment_hsv, copy_paste, letterbox, mixup, random_perspective -from utils.general import (LOGGER, check_dataset, check_requirements, check_yaml, clean_str, segments2boxes, xyn2xy, - xywh2xyxy, xywhn2xyxy, xyxy2xywhn) +from utils.general import (LOGGER, NUM_THREADS, check_dataset, check_requirements, check_yaml, clean_str, + segments2boxes, xyn2xy, xywh2xyxy, xywhn2xyxy, xyxy2xywhn) from utils.torch_utils import torch_distributed_zero_first # Parameters @@ -35,7 +35,6 @@ IMG_FORMATS = ['bmp', 'jpg', 'jpeg', 'png', 'tif', 'tiff', 'dng', 'webp', 'mpo'] # acceptable image suffixes VID_FORMATS = ['mov', 'avi', 'mp4', 'mpg', 'mpeg', 'm4v', 'wmv', 'mkv'] # acceptable video suffixes WORLD_SIZE = int(os.getenv('WORLD_SIZE', 1)) # DPP -NUM_THREADS = min(8, max(1, os.cpu_count() - 1)) # number of multiprocessing threads # Get orientation exif tag for orientation in ExifTags.TAGS.keys(): diff --git a/utils/general.py b/utils/general.py index 6cc658cc3150..1da8a147510e 100755 --- a/utils/general.py +++ b/utils/general.py @@ -33,14 +33,15 @@ from utils.metrics import box_iou, fitness # Settings +FILE = Path(__file__).resolve() +ROOT = FILE.parents[1] # YOLOv5 root directory +NUM_THREADS = min(8, max(1, os.cpu_count() - 1)) # number of YOLOv5 multiprocessing threads + torch.set_printoptions(linewidth=320, precision=5, profile='long') np.set_printoptions(linewidth=320, formatter={'float_kind': '{:11.5g}'.format}) # format short g, %precision=5 pd.options.display.max_columns = 10 cv2.setNumThreads(0) # prevent OpenCV from multithreading (incompatible with PyTorch DataLoader) -os.environ['NUMEXPR_MAX_THREADS'] = str(min(os.cpu_count(), 8)) # NumExpr max threads - -FILE = Path(__file__).resolve() -ROOT = FILE.parents[1] # YOLOv5 root directory +os.environ['NUMEXPR_MAX_THREADS'] = str(NUM_THREADS) # NumExpr max threads def set_logging(name=None, verbose=True):