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

ultralytics 8.2.13 enable --slow CUDA tests (#12511)

Browse source 
Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Glenn Jocher 2024-05-11 22:33:02 +02:00 • committed by GitHub
parent 756a224bcd
commit 274d6e7e74
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 10 additions and 12 deletions
Download patch file Download diff file Expand all files Collapse all files

9
tests/test_cuda.py
View file

@ -19,14 +19,6 @@ def test_checks():
assert torch.cuda.device_count() == CUDA_DEVICE_COUNT
@pytest.mark.slow
@pytest.mark.skipif(not CUDA_IS_AVAILABLE, reason="CUDA is not available")
def test_export_engine():
"""Test exporting the YOLO model to NVIDIA TensorRT format."""
f = YOLO(MODEL).export(format="engine", device=0)
YOLO(f)(SOURCE, device=0)
@pytest.mark.slow
@pytest.mark.skipif(not CUDA_IS_AVAILABLE, reason="CUDA is not available")
@pytest.mark.parametrize(
@ -49,6 +41,7 @@ def test_export_engine_matrix(task, dynamic, int8, half, batch):
half=half,
batch=batch,
data=TASK2DATA[task],
workspace=1, # reduce workspace GB for less resource utilization during testing
)
YOLO(file)([SOURCE] * batch, imgsz=64 if dynamic else 32) # exported model inference
Path(file).unlink() # cleanup