110 lines
3.3 KiB
Python
110 lines
3.3 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
|
|
import pytest
|
|
|
|
from yolo_webui.config import AugmentationConfig, DatasetSplitConfig, MlflowConfig, TrainingConfig
|
|
from yolo_webui.trainer import TrainingRunner, mlflow_environment
|
|
|
|
|
|
def test_train_kwargs_omit_optional_empty_values() -> None:
|
|
config = TrainingConfig(
|
|
dataset=" dataset.yaml ",
|
|
model=" model.pt ",
|
|
augmentation=AugmentationConfig(enabled=False),
|
|
)
|
|
|
|
assert config.train_kwargs() == {
|
|
"data": "dataset.yaml",
|
|
"epochs": 100,
|
|
"imgsz": 640,
|
|
"batch": 16,
|
|
"workers": 8,
|
|
"patience": 100,
|
|
"project": "runs/train",
|
|
"verbose": True,
|
|
}
|
|
|
|
|
|
def test_augmentation_kwargs_are_passed_to_ultralytics() -> None:
|
|
config = TrainingConfig(
|
|
dataset="dataset.yaml",
|
|
model="model.pt",
|
|
augmentation=AugmentationConfig(mosaic=0.5, mixup=0.2, close_mosaic=5),
|
|
)
|
|
|
|
kwargs = config.train_kwargs()
|
|
|
|
assert kwargs["mosaic"] == 0.5
|
|
assert kwargs["mixup"] == 0.2
|
|
assert kwargs["close_mosaic"] == 5
|
|
assert kwargs["auto_augment"] == "randaugment"
|
|
|
|
|
|
@pytest.mark.parametrize("field", ["mosaic", "fliplr", "erasing", "perspective"])
|
|
def test_augmentation_probabilities_are_validated(field: str) -> None:
|
|
augmentation = AugmentationConfig(**{field: 1.1})
|
|
|
|
with pytest.raises(ValueError, match="от 0 до 1"):
|
|
augmentation.validate()
|
|
|
|
|
|
@pytest.mark.parametrize("batch", [0, -2])
|
|
def test_invalid_batch_is_rejected(batch: int) -> None:
|
|
config = TrainingConfig(dataset="dataset.yaml", model="model.pt", batch_size=batch)
|
|
|
|
with pytest.raises(ValueError, match="Batch"):
|
|
config.validate()
|
|
|
|
|
|
def test_mlflow_environment_is_restored(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setenv("MLFLOW_TRACKING_URI", "previous")
|
|
config = MlflowConfig(
|
|
enabled=True,
|
|
tracking_uri="http://127.0.0.1:5000",
|
|
experiment_name="test-experiment",
|
|
run_name="test-run",
|
|
)
|
|
|
|
with mlflow_environment(config):
|
|
assert os.environ["MLFLOW_TRACKING_URI"] == config.tracking_uri
|
|
assert os.environ["MLFLOW_EXPERIMENT_NAME"] == config.experiment_name
|
|
assert os.environ["MLFLOW_RUN"] == config.run_name
|
|
|
|
assert os.environ["MLFLOW_TRACKING_URI"] == "previous"
|
|
assert "MLFLOW_EXPERIMENT_NAME" not in os.environ
|
|
assert "MLFLOW_RUN" not in os.environ
|
|
|
|
|
|
def test_metrics_summary_skips_non_numeric_values() -> None:
|
|
summary = TrainingRunner._metrics_summary(
|
|
{
|
|
"metrics/mAP50": 0.81234,
|
|
"label": "invalid",
|
|
"val/loss": 0.12345,
|
|
"metrics/precision": 0.95678,
|
|
"another": 12.34,
|
|
}
|
|
)
|
|
|
|
assert summary == "mAP50=0.8123 · loss=0.1235 · precision=0.9568"
|
|
|
|
|
|
@pytest.mark.parametrize("ratio", [0.05, 0.98])
|
|
def test_dataset_split_ratio_is_validated(ratio: float) -> None:
|
|
config = DatasetSplitConfig(enabled=True, train_ratio=ratio)
|
|
with pytest.raises(ValueError, match="Доля обучающей выборки"):
|
|
config.validate()
|
|
|
|
|
|
def test_classification_rejects_detection_style_auto_split() -> None:
|
|
config = TrainingConfig(
|
|
dataset="classification-dataset",
|
|
model="model.pt",
|
|
task="classify",
|
|
split=DatasetSplitConfig(enabled=True),
|
|
)
|
|
|
|
with pytest.raises(ValueError, match="classify"):
|
|
config.validate()
|