# 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, semantic 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), ( semantic.SemanticSegmentationTrainer, semantic.SemanticSegmentationValidator, semantic.SemanticSegmentationPredictor, "cityscapes8.yaml", "yolo26n-sem.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"