Spaces:
Sleeping
Sleeping
File size: 4,682 Bytes
339f0eb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 |
# Ultralytics YOLO 🚀, AGPL-3.0 license
import contextlib
from pathlib import Path
import pytest
from ultralytics import YOLO, download
from ultralytics.utils import ASSETS, DATASETS_DIR, ROOT, SETTINGS, WEIGHTS_DIR
from ultralytics.utils.checks import check_requirements
MODEL = WEIGHTS_DIR / 'path with spaces' / 'yolov8n.pt' # test spaces in path
CFG = 'yolov8n.yaml'
SOURCE = ASSETS / 'bus.jpg'
TMP = (ROOT / '../tests/tmp').resolve() # temp directory for test files
@pytest.mark.skipif(not check_requirements('ray', install=False), reason='ray[tune] not installed')
def test_model_ray_tune():
"""Tune YOLO model with Ray optimization library."""
YOLO('yolov8n-cls.yaml').tune(use_ray=True,
data='imagenet10',
grace_period=1,
iterations=1,
imgsz=32,
epochs=1,
plots=False,
device='cpu')
@pytest.mark.skipif(not check_requirements('mlflow', install=False), reason='mlflow not installed')
def test_mlflow():
"""Test training with MLflow tracking enabled."""
SETTINGS['mlflow'] = True
YOLO('yolov8n-cls.yaml').train(data='imagenet10', imgsz=32, epochs=3, plots=False, device='cpu')
@pytest.mark.skipif(not check_requirements('tritonclient', install=False), reason='tritonclient[all] not installed')
def test_triton():
"""Test NVIDIA Triton Server functionalities."""
check_requirements('tritonclient[all]')
import subprocess
import time
from tritonclient.http import InferenceServerClient # noqa
# Create variables
model_name = 'yolo'
triton_repo_path = TMP / 'triton_repo'
triton_model_path = triton_repo_path / model_name
# Export model to ONNX
f = YOLO(MODEL).export(format='onnx', dynamic=True)
# Prepare Triton repo
(triton_model_path / '1').mkdir(parents=True, exist_ok=True)
Path(f).rename(triton_model_path / '1' / 'model.onnx')
(triton_model_path / 'config.pdtxt').touch()
# Define image https://catalog.ngc.nvidia.com/orgs/nvidia/containers/tritonserver
tag = 'nvcr.io/nvidia/tritonserver:23.09-py3' # 6.4 GB
# Pull the image
subprocess.call(f'docker pull {tag}', shell=True)
# Run the Triton server and capture the container ID
container_id = subprocess.check_output(
f'docker run -d --rm -v {triton_repo_path}:/models -p 8000:8000 {tag} tritonserver --model-repository=/models',
shell=True).decode('utf-8').strip()
# Wait for the Triton server to start
triton_client = InferenceServerClient(url='localhost:8000', verbose=False, ssl=False)
# Wait until model is ready
for _ in range(10):
with contextlib.suppress(Exception):
assert triton_client.is_model_ready(model_name)
break
time.sleep(1)
# Check Triton inference
YOLO(f'http://localhost:8000/{model_name}', 'detect')(SOURCE) # exported model inference
# Kill and remove the container at the end of the test
subprocess.call(f'docker kill {container_id}', shell=True)
@pytest.mark.skipif(not check_requirements('pycocotools', install=False), reason='pycocotools not installed')
def test_pycocotools():
"""Validate model predictions using pycocotools."""
from ultralytics.models.yolo.detect import DetectionValidator
from ultralytics.models.yolo.pose import PoseValidator
from ultralytics.models.yolo.segment import SegmentationValidator
# Download annotations after each dataset downloads first
url = 'https://github.com/ultralytics/assets/releases/download/v0.0.0/'
args = {'model': 'yolov8n.pt', 'data': 'coco8.yaml', 'save_json': True, 'imgsz': 64}
validator = DetectionValidator(args=args)
validator()
validator.is_coco = True
download(f'{url}instances_val2017.json', dir=DATASETS_DIR / 'coco8/annotations')
_ = validator.eval_json(validator.stats)
args = {'model': 'yolov8n-seg.pt', 'data': 'coco8-seg.yaml', 'save_json': True, 'imgsz': 64}
validator = SegmentationValidator(args=args)
validator()
validator.is_coco = True
download(f'{url}instances_val2017.json', dir=DATASETS_DIR / 'coco8-seg/annotations')
_ = validator.eval_json(validator.stats)
args = {'model': 'yolov8n-pose.pt', 'data': 'coco8-pose.yaml', 'save_json': True, 'imgsz': 64}
validator = PoseValidator(args=args)
validator()
validator.is_coco = True
download(f'{url}person_keypoints_val2017.json', dir=DATASETS_DIR / 'coco8-pose/annotations')
_ = validator.eval_json(validator.stats)
|