✅ [Pass] the test for lightning train and validate

tests/test_tools/test_solver.py

@@ -1,38 +1,39 @@
 import sys
+from math import isclose
 from pathlib import Path
 
 import pytest
-from torch import allclose, tensor
+from lightning.pytorch import Trainer
+from torch.utils.data import DataLoader
 
 project_root = Path(__file__).resolve().parent.parent.parent
 sys.path.append(str(project_root))
 
 from yolo.config.config import Config
 from yolo.model.yolo import YOLO
-from yolo.tools.data_loader import StreamDataLoader, YoloDataLoader
-from yolo.tools.solver import ModelTester, ModelTrainer, ModelValidator
+from yolo.tools.data_loader import StreamDataLoader
+from yolo.tools.solver import TrainModel, ValidateModel
 from yolo.utils.bounding_box_utils import Anc2Box, Vec2Box
 
 
 @pytest.fixture
-def model_validator(validation_cfg: Config, model: YOLO, vec2box: Vec2Box, validation_progress_logger, device):
-    validator = ModelValidator(
-        validation_cfg.task, validation_cfg.dataset, model, vec2box, validation_progress_logger, device
-    )
+def model_validator(validation_cfg: Config):
+    validator = ValidateModel(validation_cfg)
     return validator
 
 
-def test_model_validator_initialization(model_validator: ModelValidator):
+def test_model_validator_initialization(solver: Trainer, model_validator: ValidateModel):
     assert isinstance(model_validator.model, YOLO)
-    assert hasattr(model_validator, "solve")
+    assert hasattr(solver, "validate")
 
 
-def test_model_validator_solve_mock_dataset(model_validator: ModelValidator, validation_dataloader: YoloDataLoader):
-    mAPs = model_validator.solve(validation_dataloader)
-    except_mAPs = {"mAP.5": tensor(0.6969), "mAP.5:.95": tensor(0.4195)}
-    assert allclose(mAPs["mAP.5"], except_mAPs["mAP.5"], rtol=0.1)
-    print(mAPs)
-    assert allclose(mAPs["mAP.5:.95"], except_mAPs["mAP.5:.95"], rtol=0.1)
+def test_model_validator_solve_mock_dataset(
+    solver: Trainer, model_validator: ValidateModel, validation_dataloader: DataLoader
+):
+    mAPs = solver.validate(model_validator, dataloaders=validation_dataloader)[0]
+    except_mAPs = {"map_50": 0.7379, "map": 0.5617}
+    assert isclose(mAPs["map_50"], except_mAPs["map_50"], abs_tol=1e-4)
+    assert isclose(mAPs["map"], except_mAPs["map"], abs_tol=1e-4)
 
 
 @pytest.fixture
@@ -63,17 +64,14 @@ def test_modelv7_tester_solve_single_image(modelv7_tester: ModelTester, file_str
 @pytest.fixture
 def model_trainer(train_cfg: Config, model: YOLO, vec2box: Vec2Box, train_progress_logger, device):
     train_cfg.task.epoch = 2
-    trainer = ModelTrainer(train_cfg, model, vec2box, train_progress_logger, device, use_ddp=False)
+    trainer = TrainModel(train_cfg)
     return trainer
 
 
-def test_model_trainer_initialization(model_trainer: ModelTrainer):
-
+def test_model_trainer_initialization(solver: Trainer, model_trainer: TrainModel):
     assert isinstance(model_trainer.model, YOLO)
-    assert hasattr(model_trainer, "solve")
-    assert model_trainer.optimizer is not None
-    assert model_trainer.scheduler is not None
-    assert model_trainer.loss_fn is not None
+    assert hasattr(solver, "fit")
+    assert solver.optimizers is not None
 
 
 # def test_model_trainer_solve_mock_dataset(model_trainer: ModelTrainer, train_dataloader: YoloDataLoader):