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

ultralytics 8.0.19 seg/det dataset warning and DDP-cls/seg fixes (#595)

Browse source 
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: 曾逸夫（Zeng Yifu） <41098760+Zengyf-CVer@users.noreply.github.com>
Co-authored-by: Laughing <61612323+Laughing-q@users.noreply.github.com>
This commit is contained in:
Glenn Jocher 2023-01-24 23:22:02 +01:00 • committed by GitHub
parent 936414c615
commit 520825c4b2
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
21 changed files with 103 additions and 63 deletions
Download patch file Download diff file Expand all files Collapse all files

11
ultralytics/yolo/utils/torch_utils.py
View file

@ -58,7 +58,7 @@ def DDP_model(model):
return DDP(model, device_ids=[LOCAL_RANK], output_device=LOCAL_RANK)
def select_device(device='', batch_size=0, newline=False):
def select_device(device='', batch=0, newline=False):
# device = None or 'cpu' or 0 or '0' or '0,1,2,3'
ver = git_describe() or ultralytics.__version__ # git commit or pip package version
s = f'Ultralytics YOLOv{ver} 🚀 Python-{platform.python_version()} torch-{torch.__version__} '
@ -71,14 +71,15 @@ def select_device(device='', batch_size=0, newline=False):
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() and torch.cuda.device_count() >= len(device.replace(',', '')), \
f"Invalid CUDA 'device={device}' requested, use 'device=cpu' or pass valid CUDA device(s)"
if not (torch.cuda.is_available() and torch.cuda.device_count() >= len(device.replace(',', ''))):
raise ValueError(f"Invalid CUDA 'device={device}' requested, use 'device=cpu' or pass valid CUDA device(s)")
if not cpu and not mps and torch.cuda.is_available(): # prefer GPU if available
devices = device.split(',') if device else '0' # range(torch.cuda.device_count()) # i.e. 0,1,6,7
n = len(devices) # device count
if n > 1 and batch_size > 0: # check batch_size is divisible by device_count
assert batch_size % n == 0, f'batch-size {batch_size} not multiple of GPU count {n}'
if n > 1 and batch > 0 and batch % n != 0: # check batch_size is divisible by device_count
raise ValueError(f'batch={batch} is not multiple of GPU count {n}.\n'
f'Try batch={batch // n} or batch={batch // n + 1}')
space = ' ' * (len(s) + 1)
for i, d in enumerate(devices):
p = torch.cuda.get_device_properties(i)