mirror of
https://github.com/ultralytics/ultralytics
synced 2026-05-06 15:12:24 +00:00
236 lines
8.6 KiB
Python
236 lines
8.6 KiB
Python
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
|
|
|
import sys
|
|
from types import SimpleNamespace
|
|
from unittest import mock
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from tests import MODEL, SOURCE, TASK_MODEL_DATA
|
|
from ultralytics import YOLO
|
|
from ultralytics.cfg import get_cfg
|
|
from ultralytics.engine.exporter import Exporter
|
|
from ultralytics.engine.trainer import BaseTrainer
|
|
from ultralytics.models.yolo import classify, detect, obb, pose, segment
|
|
from ultralytics.nn.tasks import load_checkpoint
|
|
from ultralytics.utils import ASSETS, DEFAULT_CFG, WEIGHTS_DIR
|
|
|
|
|
|
def test_func(*args, **kwargs):
|
|
"""Test function used as a callback stub to verify callback registration."""
|
|
print("callback test passed")
|
|
|
|
|
|
def test_export():
|
|
"""Test model exporting functionality by adding a callback and verifying its execution."""
|
|
exporter = Exporter()
|
|
exporter.add_callback("on_export_start", test_func)
|
|
assert test_func in exporter.callbacks["on_export_start"], "on_export_start callback not registered"
|
|
f = exporter(model=YOLO("yolo26n.yaml").model)
|
|
YOLO(f)(SOURCE) # exported model inference
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"trainer_cls,validator_cls,predictor_cls,data,model,weights",
|
|
[
|
|
(
|
|
detect.DetectionTrainer,
|
|
detect.DetectionValidator,
|
|
detect.DetectionPredictor,
|
|
"coco8.yaml",
|
|
"yolo26n.yaml",
|
|
MODEL,
|
|
),
|
|
(
|
|
segment.SegmentationTrainer,
|
|
segment.SegmentationValidator,
|
|
segment.SegmentationPredictor,
|
|
"coco8-seg.yaml",
|
|
"yolo26n-seg.yaml",
|
|
WEIGHTS_DIR / "yolo26n-seg.pt",
|
|
),
|
|
(
|
|
classify.ClassificationTrainer,
|
|
classify.ClassificationValidator,
|
|
classify.ClassificationPredictor,
|
|
"imagenet10",
|
|
"yolo26n-cls.yaml",
|
|
None,
|
|
),
|
|
(obb.OBBTrainer, obb.OBBValidator, obb.OBBPredictor, "dota8.yaml", "yolo26n-obb.yaml", None),
|
|
(pose.PoseTrainer, pose.PoseValidator, pose.PosePredictor, "coco8-pose.yaml", "yolo26n-pose.yaml", None),
|
|
],
|
|
)
|
|
def test_task(trainer_cls, validator_cls, predictor_cls, data, model, weights):
|
|
"""Test YOLO training, validation, and prediction for various tasks."""
|
|
overrides = {
|
|
"data": data,
|
|
"model": model,
|
|
"imgsz": 32,
|
|
"epochs": 1,
|
|
"save": False,
|
|
"mask_ratio": 1,
|
|
"overlap_mask": False,
|
|
}
|
|
|
|
# Trainer
|
|
trainer = trainer_cls(overrides=overrides)
|
|
trainer.add_callback("on_train_start", test_func)
|
|
assert test_func in trainer.callbacks["on_train_start"], "on_train_start callback not registered"
|
|
trainer.train()
|
|
|
|
# Validator
|
|
cfg = get_cfg(DEFAULT_CFG)
|
|
cfg.data = data
|
|
cfg.imgsz = 32
|
|
val = validator_cls(args=cfg)
|
|
val.add_callback("on_val_start", test_func)
|
|
assert test_func in val.callbacks["on_val_start"], "on_val_start callback not registered"
|
|
val(model=trainer.best)
|
|
|
|
# Predictor
|
|
pred = predictor_cls(overrides={"imgsz": [64, 64]})
|
|
pred.add_callback("on_predict_start", test_func)
|
|
assert test_func in pred.callbacks["on_predict_start"], "on_predict_start callback not registered"
|
|
|
|
# Determine model path for prediction
|
|
model_path = weights if weights else trainer.best
|
|
if model == "yolo26n.yaml": # only for detection
|
|
# Confirm there is no issue with sys.argv being empty
|
|
with mock.patch.object(sys, "argv", []):
|
|
result = pred(source=ASSETS, model=model_path)
|
|
assert len(result) > 0, f"Predictor returned no results for {model}"
|
|
else:
|
|
result = pred(source=ASSETS, model=model_path)
|
|
assert len(result) > 0, f"Predictor returned no results for {model}"
|
|
|
|
# Test resume functionality
|
|
with pytest.raises(AssertionError):
|
|
trainer_cls(overrides={**overrides, "resume": trainer.last}).train()
|
|
|
|
|
|
@pytest.mark.parametrize("task,weight,data", TASK_MODEL_DATA)
|
|
def test_resume_incomplete(task, weight, data, tmp_path):
|
|
"""Test training resumes from an incomplete checkpoint."""
|
|
train_args = {
|
|
"data": data,
|
|
"epochs": 2,
|
|
"save": True,
|
|
"plots": False,
|
|
"workers": 0,
|
|
"project": tmp_path,
|
|
"name": task,
|
|
"imgsz": 32,
|
|
"exist_ok": True,
|
|
}
|
|
|
|
def stop_after_first_epoch(trainer):
|
|
if trainer.epoch == 0:
|
|
trainer.stop = True
|
|
|
|
def disable_final_eval(trainer):
|
|
trainer.final_eval = lambda: None
|
|
|
|
model = YOLO(weight)
|
|
model.add_callback("on_train_start", disable_final_eval)
|
|
model.add_callback("on_train_epoch_end", stop_after_first_epoch)
|
|
model.train(**train_args)
|
|
last_path = model.trainer.last
|
|
_, ckpt = load_checkpoint(last_path)
|
|
assert ckpt["epoch"] == 0, "checkpoint should be resumable"
|
|
|
|
# Resume training using the checkpoint
|
|
resume_model = YOLO(last_path)
|
|
resume_model.train(resume=True, **train_args)
|
|
assert resume_model.trainer.start_epoch == resume_model.trainer.epoch == 1, "resume test failed"
|
|
|
|
|
|
def test_nan_recovery():
|
|
"""Test NaN loss detection and recovery during training."""
|
|
nan_injected = [False]
|
|
|
|
def inject_nan(trainer):
|
|
"""Inject NaN into loss during batch processing to test recovery mechanism."""
|
|
if trainer.epoch == 1 and trainer.tloss is not None and not nan_injected[0]:
|
|
trainer.tloss *= torch.tensor(float("nan"))
|
|
nan_injected[0] = True
|
|
|
|
overrides = {"data": "coco8.yaml", "model": "yolo26n.yaml", "imgsz": 32, "epochs": 3}
|
|
trainer = detect.DetectionTrainer(overrides=overrides)
|
|
trainer.add_callback("on_train_batch_end", inject_nan)
|
|
trainer.train()
|
|
assert nan_injected[0], "NaN injection failed"
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"kwargs,uses_weights",
|
|
[({}, True), ({"pretrained": True}, True), ({"pretrained": False}, False), ({"pretrained": MODEL}, True)],
|
|
)
|
|
def test_train_reuses_loaded_checkpoint_model(monkeypatch, kwargs, uses_weights):
|
|
"""Test training reuses loaded checkpoint config while respecting the pretrained argument."""
|
|
model = YOLO("yolo26n.yaml")
|
|
model.ckpt = {"checkpoint": True}
|
|
model.ckpt_path = "/tmp/fake.pt"
|
|
model.overrides["model"] = "ul://glenn-jocher/m2/exp-14"
|
|
model.overrides["pretrained"] = False
|
|
original_model = model.model
|
|
captured = {}
|
|
|
|
class FakeTrainer:
|
|
def __init__(self, overrides=None, _callbacks=None):
|
|
self.overrides = overrides
|
|
self.callbacks = _callbacks
|
|
self.model = None
|
|
self.validator = SimpleNamespace(metrics=None)
|
|
self.best = MODEL.parent / "nonexistent-best.pt"
|
|
self.last = MODEL
|
|
captured["trainer"] = self
|
|
|
|
def get_model(self, cfg=None, weights=None, verbose=True):
|
|
captured["cfg"] = cfg
|
|
captured["weights"] = weights
|
|
return original_model
|
|
|
|
def train(self):
|
|
return None
|
|
|
|
monkeypatch.setattr("ultralytics.engine.model.checks.check_pip_update_available", lambda: None)
|
|
monkeypatch.setattr(model, "_smart_load", lambda key: FakeTrainer)
|
|
monkeypatch.setattr(
|
|
"ultralytics.engine.model.load_checkpoint",
|
|
lambda path: (original_model, {"checkpoint": True}),
|
|
)
|
|
|
|
model.train(data="coco8.yaml", epochs=1, **kwargs)
|
|
|
|
assert captured["trainer"].model is original_model, "Trainer model does not match original"
|
|
assert captured["cfg"] == original_model.yaml, f"Config mismatch: {captured['cfg']} != {original_model.yaml}"
|
|
assert captured["weights"] is (original_model if uses_weights else None), "Unexpected weights loaded"
|
|
|
|
|
|
@pytest.mark.parametrize("pretrained,uses_weights", [(True, True), (False, False), (MODEL, True)])
|
|
def test_setup_model_respects_pretrained_arg_for_pt_models(monkeypatch, pretrained, uses_weights):
|
|
"""Test .pt models use checkpoint config while respecting the pretrained argument."""
|
|
captured = {}
|
|
checkpoint_model = SimpleNamespace(yaml={"nc": 80})
|
|
trainer = object.__new__(BaseTrainer)
|
|
trainer.model = "yolo26n.pt"
|
|
trainer.args = SimpleNamespace(pretrained=pretrained)
|
|
trainer.resume = False
|
|
|
|
def fake_get_model(cfg=None, weights=None, verbose=True):
|
|
captured["cfg"] = cfg
|
|
captured["weights"] = weights
|
|
return SimpleNamespace()
|
|
|
|
trainer.get_model = fake_get_model
|
|
monkeypatch.setattr(
|
|
"ultralytics.engine.trainer.load_checkpoint", lambda path: (checkpoint_model, {"checkpoint": True})
|
|
)
|
|
|
|
trainer.setup_model()
|
|
|
|
assert captured["cfg"] == checkpoint_model.yaml, "Checkpoint config was not used"
|
|
assert captured["weights"] is (checkpoint_model if uses_weights else None), "Unexpected weights loaded"
|