Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Kalen Michael <kalenmike@gmail.com>
This commit is contained in:
Glenn Jocher 2023-01-09 23:22:33 +01:00 committed by GitHub
parent 6feba17760
commit 422c49d439
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
97 changed files with 224 additions and 757 deletions

View file

@ -1,3 +1,5 @@
# Ultralytics YOLO 🚀, GPL-3.0 license
import math
import os
import platform
@ -59,7 +61,7 @@ def DDP_model(model):
def select_device(device='', batch_size=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 YOLO 🚀 {ver} Python-{platform.python_version()} torch-{torch.__version__} '
s = f'Ultralytics YOLOv{ver} 🚀 Python-{platform.python_version()} torch-{torch.__version__} '
device = str(device).strip().lower().replace('cuda:', '').replace('none', '') # to string, 'cuda:0' to '0'
cpu = device == 'cpu'
mps = device == 'mps' # Apple Metal Performance Shaders (MPS)