From 58c404d8d81ebdd810954ffeead137853e98b322 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Fri, 28 Jan 2022 07:53:59 +0100 Subject: [PATCH 1/6] Fix2 select_device() for Multi-GPU --- utils/datasets.py | 6 +++--- utils/torch_utils.py | 11 +++++++---- 2 files changed, 10 insertions(+), 7 deletions(-) diff --git a/utils/datasets.py b/utils/datasets.py index 4eb444087860..07f6321e0285 100755 --- a/utils/datasets.py +++ b/utils/datasets.py @@ -29,13 +29,12 @@ from utils.augmentations import Albumentations, augment_hsv, copy_paste, letterbox, mixup, random_perspective 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 device_count, torch_distributed_zero_first +from utils.torch_utils import torch_distributed_zero_first # Parameters HELP_URL = 'https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data' IMG_FORMATS = ['bmp', 'dng', 'jpeg', 'jpg', 'mpo', 'png', 'tif', 'tiff', 'webp'] # include image suffixes VID_FORMATS = ['asf', 'avi', 'gif', 'm4v', 'mkv', 'mov', 'mp4', 'mpeg', 'mpg', 'wmv'] # include video suffixes -DEVICE_COUNT = max(device_count(), 1) # number of CUDA devices # Get orientation exif tag for orientation in ExifTags.TAGS.keys(): @@ -110,7 +109,8 @@ def create_dataloader(path, imgsz, batch_size, stride, single_cls=False, hyp=Non prefix=prefix) batch_size = min(batch_size, len(dataset)) - nw = min([os.cpu_count() // DEVICE_COUNT, batch_size if batch_size > 1 else 0, workers]) # number of workers + nd = torch.cuda.device_count() # number of CUDA devices + nw = min([os.cpu_count() // max(nd, 1), batch_size if batch_size > 1 else 0, workers]) # number of workers sampler = None if rank == -1 else distributed.DistributedSampler(dataset, shuffle=shuffle) loader = DataLoader if image_weights else InfiniteDataLoader # only DataLoader allows for attribute updates return loader(dataset, diff --git a/utils/torch_utils.py b/utils/torch_utils.py index d958a8951074..552ed65e9578 100644 --- a/utils/torch_utils.py +++ b/utils/torch_utils.py @@ -8,6 +8,7 @@ import os import platform import subprocess +import sys import time from contextlib import contextmanager from copy import deepcopy @@ -54,7 +55,8 @@ def git_describe(path=Path(__file__).parent): # path must be a directory def device_count(): - # Returns number of CUDA devices available. Safe version of torch.cuda.device_count(). + # Returns number of CUDA devices available. Safe version of torch.cuda.device_count(). Only works on Linux. + assert sys.platform == 'Linux' try: cmd = 'nvidia-smi -L | wc -l' return int(subprocess.run(cmd, shell=True, capture_output=True, check=True).stdout.decode().split()[-1]) @@ -63,6 +65,7 @@ def device_count(): def select_device(device='', batch_size=0, newline=True): + print(device) # device = 'cpu' or '0' or '0,1,2,3' s = f'YOLOv5 🚀 {git_describe() or date_modified()} torch {torch.__version__} ' # string device = str(device).strip().lower().replace('cuda:', '') # to string, 'cuda:0' to '0' @@ -70,10 +73,10 @@ def select_device(device='', batch_size=0, newline=True): if cpu: os.environ['CUDA_VISIBLE_DEVICES'] = '-1' # force torch.cuda.is_available() = False elif device: # non-cpu device requested - nd = device_count() # number of CUDA devices - assert nd > int(max(device.split(','))), f'Invalid `--device {device}` request, valid devices are 0 - {nd - 1}' os.environ['CUDA_VISIBLE_DEVICES'] = device # set environment variable - must be before assert is_available() - assert torch.cuda.is_available(), 'CUDA is not available, use `--device cpu` or do not pass a --device' + nd = torch.cuda.device_count() # number of CUDA devices + assert torch.cuda.is_available(), f'CUDA `--device {device}` requested but `torch.cuda.is_available()=False`' + assert nd > 0, f'Invalid CUDA `--device {device}` requested, use `--device cpu` or pass valid CUDA device(s)' cuda = not cpu and torch.cuda.is_available() if cuda: From 881411d5abb965ad39b439a8ce9510443d885b31 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Fri, 28 Jan 2022 07:56:24 +0100 Subject: [PATCH 2/6] Cleanup --- utils/torch_utils.py | 1 - 1 file changed, 1 deletion(-) diff --git a/utils/torch_utils.py b/utils/torch_utils.py index 552ed65e9578..6227b103f2fc 100644 --- a/utils/torch_utils.py +++ b/utils/torch_utils.py @@ -65,7 +65,6 @@ def device_count(): def select_device(device='', batch_size=0, newline=True): - print(device) # device = 'cpu' or '0' or '0,1,2,3' s = f'YOLOv5 🚀 {git_describe() or date_modified()} torch {torch.__version__} ' # string device = str(device).strip().lower().replace('cuda:', '') # to string, 'cuda:0' to '0' From 547d373b4d6b0093f80e9d3ed5be9b60c3915650 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Fri, 28 Jan 2022 07:58:24 +0100 Subject: [PATCH 3/6] Cleanup --- utils/torch_utils.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/utils/torch_utils.py b/utils/torch_utils.py index 6227b103f2fc..e1b61b1b5feb 100644 --- a/utils/torch_utils.py +++ b/utils/torch_utils.py @@ -8,7 +8,6 @@ import os import platform import subprocess -import sys import time from contextlib import contextmanager from copy import deepcopy @@ -56,7 +55,7 @@ def git_describe(path=Path(__file__).parent): # path must be a directory def device_count(): # Returns number of CUDA devices available. Safe version of torch.cuda.device_count(). Only works on Linux. - assert sys.platform == 'Linux' + assert platform.system() == 'Linux' try: cmd = 'nvidia-smi -L | wc -l' return int(subprocess.run(cmd, shell=True, capture_output=True, check=True).stdout.decode().split()[-1]) From 194ea1abb4d56c3eace0b4bbd39e83bcd17b59a0 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Fri, 28 Jan 2022 08:02:23 +0100 Subject: [PATCH 4/6] Simplify error message --- utils/torch_utils.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/utils/torch_utils.py b/utils/torch_utils.py index e1b61b1b5feb..0fd9f533ed66 100644 --- a/utils/torch_utils.py +++ b/utils/torch_utils.py @@ -72,9 +72,8 @@ def select_device(device='', batch_size=0, newline=True): os.environ['CUDA_VISIBLE_DEVICES'] = '-1' # force torch.cuda.is_available() = False elif device: # non-cpu device requested os.environ['CUDA_VISIBLE_DEVICES'] = device # set environment variable - must be before assert is_available() - nd = torch.cuda.device_count() # number of CUDA devices - assert torch.cuda.is_available(), f'CUDA `--device {device}` requested but `torch.cuda.is_available()=False`' - assert nd > 0, f'Invalid CUDA `--device {device}` requested, use `--device cpu` or pass valid CUDA device(s)' + assert torch.cuda.is_available(), \ + f"Invalid CUDA '--device {device}' requested, use '--device cpu' or pass valid CUDA device(s)" cuda = not cpu and torch.cuda.is_available() if cuda: From 36bb7eda33cd392b8720dc7badc110ad7e6a3b99 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Fri, 28 Jan 2022 08:04:47 +0100 Subject: [PATCH 5/6] Improve assert --- utils/torch_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/utils/torch_utils.py b/utils/torch_utils.py index 0fd9f533ed66..941ec95afa35 100644 --- a/utils/torch_utils.py +++ b/utils/torch_utils.py @@ -72,7 +72,7 @@ def select_device(device='', batch_size=0, newline=True): os.environ['CUDA_VISIBLE_DEVICES'] = '-1' # force torch.cuda.is_available() = False elif device: # non-cpu device requested os.environ['CUDA_VISIBLE_DEVICES'] = device # set environment variable - must be before assert is_available() - assert torch.cuda.is_available(), \ + assert torch.cuda.is_available() and torch.cuda.device_count() >= len(device.replace(',', '')), \ f"Invalid CUDA '--device {device}' requested, use '--device cpu' or pass valid CUDA device(s)" cuda = not cpu and torch.cuda.is_available() From 09980a6a20ed5f88f76ee7843e71eda637e93825 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Fri, 28 Jan 2022 08:17:06 +0100 Subject: [PATCH 6/6] Update torch_utils.py --- utils/torch_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/utils/torch_utils.py b/utils/torch_utils.py index 941ec95afa35..2b51821a3b62 100644 --- a/utils/torch_utils.py +++ b/utils/torch_utils.py @@ -55,7 +55,7 @@ def git_describe(path=Path(__file__).parent): # path must be a directory def device_count(): # Returns number of CUDA devices available. Safe version of torch.cuda.device_count(). Only works on Linux. - assert platform.system() == 'Linux' + assert platform.system() == 'Linux', 'device_count() function only works on Linux' try: cmd = 'nvidia-smi -L | wc -l' return int(subprocess.run(cmd, shell=True, capture_output=True, check=True).stdout.decode().split()[-1])