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

Add new @Retry() decorator (#7854)

Browse source 
Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Glenn Jocher 2024-01-27 20:07:31 +01:00 • committed by GitHub
parent 5f00fbd227
commit 1435f0e9de
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 372 additions and 269 deletions
Download patch file Download diff file Expand all files Collapse all files

46
tests/test_cuda.py
View file

@ -9,9 +9,9 @@ from ultralytics.utils import ASSETS, WEIGHTS_DIR, checks
CUDA_IS_AVAILABLE = checks.cuda_is_available()
CUDA_DEVICE_COUNT = checks.cuda_device_count()
MODEL = WEIGHTS_DIR / 'path with spaces' / 'yolov8n.pt' # test spaces in path
DATA = 'coco8.yaml'
BUS = ASSETS / 'bus.jpg'
MODEL = WEIGHTS_DIR / "path with spaces" / "yolov8n.pt" # test spaces in path
DATA = "coco8.yaml"
BUS = ASSETS / "bus.jpg"
def test_checks():
@ -20,7 +20,7 @@ def test_checks():
assert torch.cuda.device_count() == CUDA_DEVICE_COUNT
@pytest.mark.skipif(not CUDA_IS_AVAILABLE, reason='CUDA is not available')
@pytest.mark.skipif(not CUDA_IS_AVAILABLE, reason="CUDA is not available")
def test_train():
"""Test model training on a minimal dataset."""
device = 0 if CUDA_DEVICE_COUNT == 1 else [0, 1]
@ -28,32 +28,32 @@ def test_train():
@pytest.mark.slow
@pytest.mark.skipif(not CUDA_IS_AVAILABLE, reason='CUDA is not available')
@pytest.mark.skipif(not CUDA_IS_AVAILABLE, reason="CUDA is not available")
def test_predict_multiple_devices():
"""Validate model prediction on multiple devices."""
model = YOLO('yolov8n.pt')
model = YOLO("yolov8n.pt")
model = model.cpu()
assert str(model.device) == 'cpu'
assert str(model.device) == "cpu"
_ = model(BUS) # CPU inference
assert str(model.device) == 'cpu'
assert str(model.device) == "cpu"
model = model.to('cuda:0')
assert str(model.device) == 'cuda:0'
model = model.to("cuda:0")
assert str(model.device) == "cuda:0"
_ = model(BUS) # CUDA inference
assert str(model.device) == 'cuda:0'
assert str(model.device) == "cuda:0"
model = model.cpu()
assert str(model.device) == 'cpu'
assert str(model.device) == "cpu"
_ = model(BUS) # CPU inference
assert str(model.device) == 'cpu'
assert str(model.device) == "cpu"
model = model.cuda()
assert str(model.device) == 'cuda:0'
assert str(model.device) == "cuda:0"
_ = model(BUS) # CUDA inference
assert str(model.device) == 'cuda:0'
assert str(model.device) == "cuda:0"
@pytest.mark.skipif(not CUDA_IS_AVAILABLE, reason='CUDA is not available')
@pytest.mark.skipif(not CUDA_IS_AVAILABLE, reason="CUDA is not available")
def test_autobatch():
"""Check batch size for YOLO model using autobatch."""
from ultralytics.utils.autobatch import check_train_batch_size
@ -62,24 +62,24 @@ def test_autobatch():
@pytest.mark.slow
@pytest.mark.skipif(not CUDA_IS_AVAILABLE, reason='CUDA is not available')
@pytest.mark.skipif(not CUDA_IS_AVAILABLE, reason="CUDA is not available")
def test_utils_benchmarks():
"""Profile YOLO models for performance benchmarks."""
from ultralytics.utils.benchmarks import ProfileModels
# Pre-export a dynamic engine model to use dynamic inference
YOLO(MODEL).export(format='engine', imgsz=32, dynamic=True, batch=1)
YOLO(MODEL).export(format="engine", imgsz=32, dynamic=True, batch=1)
ProfileModels([MODEL], imgsz=32, half=False, min_time=1, num_timed_runs=3, num_warmup_runs=1).profile()
@pytest.mark.skipif(not CUDA_IS_AVAILABLE, reason='CUDA is not available')
@pytest.mark.skipif(not CUDA_IS_AVAILABLE, reason="CUDA is not available")
def test_predict_sam():
"""Test SAM model prediction with various prompts."""
from ultralytics import SAM
from ultralytics.models.sam import Predictor as SAMPredictor
# Load a model
model = SAM(WEIGHTS_DIR / 'sam_b.pt')
model = SAM(WEIGHTS_DIR / "sam_b.pt")
# Display model information (optional)
model.info()
@ -91,14 +91,14 @@ def test_predict_sam():
model(BUS, bboxes=[439, 437, 524, 709], device=0)
# Run inference with points prompt
model(ASSETS / 'zidane.jpg', points=[900, 370], labels=[1], device=0)
model(ASSETS / "zidane.jpg", points=[900, 370], labels=[1], device=0)
# Create SAMPredictor
overrides = dict(conf=0.25, task='segment', mode='predict', imgsz=1024, model=WEIGHTS_DIR / 'mobile_sam.pt')
overrides = dict(conf=0.25, task="segment", mode="predict", imgsz=1024, model=WEIGHTS_DIR / "mobile_sam.pt")
predictor = SAMPredictor(overrides=overrides)
# Set image
predictor.set_image(ASSETS / 'zidane.jpg') # set with image file
predictor.set_image(ASSETS / "zidane.jpg") # set with image file
# predictor(bboxes=[439, 437, 524, 709])
# predictor(points=[900, 370], labels=[1])