Skip to content

Commit

Permalink
Refactor NUM_THREADS (#5954)
Browse files Browse the repository at this point in the history
  • Loading branch information
glenn-jocher authored Dec 11, 2021
1 parent 2c63175 commit 8f875d9
Show file tree
Hide file tree
Showing 2 changed files with 7 additions and 7 deletions.
5 changes: 2 additions & 3 deletions utils/datasets.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,16 +26,15 @@
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
HELP_URL = 'https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data'
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():
Expand Down
9 changes: 5 additions & 4 deletions utils/general.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down

0 comments on commit 8f875d9

Please sign in to comment.