Register
Login
Resources
Docs Blog Datasets Glossary Case Studies Tutorials & Webinars
Product
Data Engine LLMs Platform Enterprise
Pricing Explore
Connect to our Discord channel

#848 Feature/sg 812 return dictionary from test

Merged
Ghost merged 1 commits into Deci-AI:master from deci-ai:feature/SG-812_return_dictionary_from_test
@@ -146,8 +146,8 @@ test_data_loader = cifar10_val()
 7. Launch test:
 7. Launch test:
 
 
 ```python
 ```python
-accuracy, top5 = trainer.test(model=model, test_loader=test_data_loader, test_metrics_list=test_metrics)
-print(f"Test results: Accuracy: {accuracy}, Top5: {top5}")
+test_results = trainer.test(model=model, test_loader=test_data_loader, test_metrics_list=test_metrics)
+print(f"Test results: Accuracy: {test_results['Accuracy']}, Top5: {test_results['Top5']}")
 ```
 ```
 
 
 </details>
 </details>
Discard
@@ -10,7 +10,6 @@ from super_gradients.training import MultiGPUMode
 from super_gradients.training import models as sg_models
 from super_gradients.training import models as sg_models
 from super_gradients.training.dataloaders import imagenet_train, imagenet_val
 from super_gradients.training.dataloaders import imagenet_train, imagenet_val
 from super_gradients.training.metrics import Accuracy, Top5
 from super_gradients.training.metrics import Accuracy, Top5
-from super_gradients.training.metrics.metric_utils import get_metrics_dict
 from super_gradients.training.models.classification_models.resnet import Bottleneck
 from super_gradients.training.models.classification_models.resnet import Bottleneck
 from super_gradients.training.models.classification_models.resnet import Bottleneck as sg_Bottleneck
 from super_gradients.training.models.classification_models.resnet import Bottleneck as sg_Bottleneck
 from super_gradients.training.utils.quantization.calibrator import QuantizationCalibrator
 from super_gradients.training.utils.quantization.calibrator import QuantizationCalibrator
@@ -112,8 +111,7 @@ if __name__ == "__main__":
 
 
     trainer.train(model=model, training_params=train_params, train_loader=train_dataloader, valid_loader=val_dataloader)
     trainer.train(model=model, training_params=train_params, train_loader=train_dataloader, valid_loader=val_dataloader)
 
 
-    val_results_tuple = trainer.test(model=model, test_loader=val_dataloader, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)
-    valid_metrics_dict = get_metrics_dict(val_results_tuple, trainer.test_metrics, trainer.loss_logging_items_names)
+    valid_metrics_dict = trainer.test(model=model, test_loader=val_dataloader, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)
 
 
     export_quantized_module_to_onnx(model=model, onnx_filename=f"{args.model_name}.onnx", input_shape=(args.batch, 3, 224, 224))
     export_quantized_module_to_onnx(model=model, onnx_filename=f"{args.model_name}.onnx", input_shape=(args.batch, 3, 224, 224))
 
 
Discard
@@ -11,7 +11,6 @@ from torch import nn
 from super_gradients.common.abstractions.abstract_logger import get_logger
 from super_gradients.common.abstractions.abstract_logger import get_logger
 from super_gradients.common.environment.device_utils import device_config
 from super_gradients.common.environment.device_utils import device_config
 from super_gradients.training import utils as core_utils, models, dataloaders
 from super_gradients.training import utils as core_utils, models, dataloaders
-from super_gradients.training.metrics.metric_utils import get_metrics_dict
 from super_gradients.training.sg_trainer import Trainer
 from super_gradients.training.sg_trainer import Trainer
 from super_gradients.training.utils import get_param
 from super_gradients.training.utils import get_param
 from super_gradients.training.utils.distributed_training_utils import setup_device
 from super_gradients.training.utils.distributed_training_utils import setup_device
@@ -158,8 +157,7 @@ class QATTrainer(Trainer):
         # VALIDATE PTQ MODEL AND PRINT SUMMARY
         # VALIDATE PTQ MODEL AND PRINT SUMMARY
         logger.info("Validating PTQ model...")
         logger.info("Validating PTQ model...")
         trainer = Trainer(experiment_name=cfg.experiment_name, ckpt_root_dir=get_param(cfg, "ckpt_root_dir", default_val=None))
         trainer = Trainer(experiment_name=cfg.experiment_name, ckpt_root_dir=get_param(cfg, "ckpt_root_dir", default_val=None))
-        val_results_tuple = trainer.test(model=model, test_loader=val_dataloader, test_metrics_list=cfg.training_hyperparams.valid_metrics_list)
-        valid_metrics_dict = get_metrics_dict(val_results_tuple, trainer.test_metrics, trainer.loss_logging_items_names)
+        valid_metrics_dict = trainer.test(model=model, test_loader=val_dataloader, test_metrics_list=cfg.training_hyperparams.valid_metrics_list)
         results = ["PTQ Model Validation Results"]
         results = ["PTQ Model Validation Results"]
         results += [f"   - {metric:10}: {value}" for metric, value in valid_metrics_dict.items()]
         results += [f"   - {metric:10}: {value}" for metric, value in valid_metrics_dict.items()]
         logger.info("\n".join(results))
         logger.info("\n".join(results))
Discard
@@ -329,14 +329,13 @@ class Trainer:
         )
         )
 
 
         # TEST
         # TEST
-        val_results_tuple = trainer.test(model=model, test_loader=val_dataloader, test_metrics_list=cfg.training_hyperparams.valid_metrics_list)
+        valid_metrics_dict = trainer.test(model=model, test_loader=val_dataloader, test_metrics_list=cfg.training_hyperparams.valid_metrics_list)
 
 
-        valid_metrics_dict = get_metrics_dict(val_results_tuple, trainer.test_metrics, trainer.loss_logging_items_names)
         results = ["Validate Results"]
         results = ["Validate Results"]
         results += [f"   - {metric:10}: {value}" for metric, value in valid_metrics_dict.items()]
         results += [f"   - {metric:10}: {value}" for metric, value in valid_metrics_dict.items()]
         logger.info("\n".join(results))
         logger.info("\n".join(results))
 
 
-        return model, val_results_tuple
+        return model, valid_metrics_dict
 
 
     @classmethod
     @classmethod
     def evaluate_checkpoint(cls, experiment_name: str, ckpt_name: str = "ckpt_latest.pth", ckpt_root_dir: str = None) -> None:
     def evaluate_checkpoint(cls, experiment_name: str, ckpt_name: str = "ckpt_latest.pth", ckpt_root_dir: str = None) -> None:
@@ -1748,6 +1747,8 @@ class Trainer:
 
 
         self._first_backward = True
         self._first_backward = True
 
 
+        test_results = get_metrics_dict(test_results, self.test_metrics, self.loss_logging_items_names)
+
         return test_results
         return test_results
 
 
     def _validate_epoch(self, epoch: int, silent_mode: bool = False) -> tuple:
     def _validate_epoch(self, epoch: int, silent_mode: bool = False) -> tuple:
Discard
@@ -242,8 +242,8 @@ class PretrainedModelsTest(unittest.TestCase):
     def test_pretrained_resnet50_imagenet(self):
     def test_pretrained_resnet50_imagenet(self):
         trainer = Trainer("imagenet_pretrained_resnet50")
         trainer = Trainer("imagenet_pretrained_resnet50")
         model = models.get(Models.RESNET50, arch_params=self.imagenet_pretrained_arch_params["resnet"], **self.imagenet_pretrained_ckpt_params)
         model = models.get(Models.RESNET50, arch_params=self.imagenet_pretrained_arch_params["resnet"], **self.imagenet_pretrained_ckpt_params)
-        res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)[0].cpu().item()
-        self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies[Models.RESNET50], delta=0.001)
+        res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)
+        self.assertAlmostEqual(res["Accuracy"].cpu().item(), self.imagenet_pretrained_accuracies[Models.RESNET50], delta=0.001)
 
 
     def test_transfer_learning_resnet50_imagenet(self):
     def test_transfer_learning_resnet50_imagenet(self):
         trainer = Trainer("imagenet_pretrained_resnet50_transfer_learning")
         trainer = Trainer("imagenet_pretrained_resnet50_transfer_learning")
@@ -259,8 +259,8 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer = Trainer("imagenet_pretrained_resnet34")
         trainer = Trainer("imagenet_pretrained_resnet34")
 
 
         model = models.get(Models.RESNET34, arch_params=self.imagenet_pretrained_arch_params["resnet"], **self.imagenet_pretrained_ckpt_params)
         model = models.get(Models.RESNET34, arch_params=self.imagenet_pretrained_arch_params["resnet"], **self.imagenet_pretrained_ckpt_params)
-        res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)[0].cpu().item()
-        self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies[Models.RESNET34], delta=0.001)
+        res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)
+        self.assertAlmostEqual(res["Accuracy"].cpu().item(), self.imagenet_pretrained_accuracies[Models.RESNET34], delta=0.001)
 
 
     def test_transfer_learning_resnet34_imagenet(self):
     def test_transfer_learning_resnet34_imagenet(self):
         trainer = Trainer("imagenet_pretrained_resnet34_transfer_learning")
         trainer = Trainer("imagenet_pretrained_resnet34_transfer_learning")
@@ -276,8 +276,8 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer = Trainer("imagenet_pretrained_resnet18")
         trainer = Trainer("imagenet_pretrained_resnet18")
 
 
         model = models.get(Models.RESNET18, arch_params=self.imagenet_pretrained_arch_params["resnet"], **self.imagenet_pretrained_ckpt_params)
         model = models.get(Models.RESNET18, arch_params=self.imagenet_pretrained_arch_params["resnet"], **self.imagenet_pretrained_ckpt_params)
-        res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)[0].cpu().item()
-        self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies[Models.RESNET18], delta=0.001)
+        res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)
+        self.assertAlmostEqual(res["Accuracy"].cpu().item(), self.imagenet_pretrained_accuracies[Models.RESNET18], delta=0.001)
 
 
     def test_transfer_learning_resnet18_imagenet(self):
     def test_transfer_learning_resnet18_imagenet(self):
         trainer = Trainer("imagenet_pretrained_resnet18_transfer_learning")
         trainer = Trainer("imagenet_pretrained_resnet18_transfer_learning")
@@ -293,8 +293,8 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer = Trainer("imagenet_pretrained_regnetY800")
         trainer = Trainer("imagenet_pretrained_regnetY800")
 
 
         model = models.get(Models.REGNETY800, arch_params=self.imagenet_pretrained_arch_params["regnet"], **self.imagenet_pretrained_ckpt_params)
         model = models.get(Models.REGNETY800, arch_params=self.imagenet_pretrained_arch_params["regnet"], **self.imagenet_pretrained_ckpt_params)
-        res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)[0].cpu().item()
-        self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies[Models.REGNETY800], delta=0.001)
+        res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)
+        self.assertAlmostEqual(res["Accuracy"].cpu().item(), self.imagenet_pretrained_accuracies[Models.REGNETY800], delta=0.001)
 
 
     def test_transfer_learning_regnetY800_imagenet(self):
     def test_transfer_learning_regnetY800_imagenet(self):
         trainer = Trainer("imagenet_pretrained_regnetY800_transfer_learning")
         trainer = Trainer("imagenet_pretrained_regnetY800_transfer_learning")
@@ -310,8 +310,8 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer = Trainer("imagenet_pretrained_regnetY600")
         trainer = Trainer("imagenet_pretrained_regnetY600")
 
 
         model = models.get(Models.REGNETY600, arch_params=self.imagenet_pretrained_arch_params["regnet"], **self.imagenet_pretrained_ckpt_params)
         model = models.get(Models.REGNETY600, arch_params=self.imagenet_pretrained_arch_params["regnet"], **self.imagenet_pretrained_ckpt_params)
-        res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)[0].cpu().item()
-        self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies[Models.REGNETY600], delta=0.001)
+        res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)
+        self.assertAlmostEqual(res["Accuracy"].cpu().item(), self.imagenet_pretrained_accuracies[Models.REGNETY600], delta=0.001)
 
 
     def test_transfer_learning_regnetY600_imagenet(self):
     def test_transfer_learning_regnetY600_imagenet(self):
         trainer = Trainer("imagenet_pretrained_regnetY600_transfer_learning")
         trainer = Trainer("imagenet_pretrained_regnetY600_transfer_learning")
@@ -327,8 +327,8 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer = Trainer("imagenet_pretrained_regnetY400")
         trainer = Trainer("imagenet_pretrained_regnetY400")
 
 
         model = models.get(Models.REGNETY400, arch_params=self.imagenet_pretrained_arch_params["regnet"], **self.imagenet_pretrained_ckpt_params)
         model = models.get(Models.REGNETY400, arch_params=self.imagenet_pretrained_arch_params["regnet"], **self.imagenet_pretrained_ckpt_params)
-        res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)[0].cpu().item()
-        self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies[Models.REGNETY400], delta=0.001)
+        res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)
+        self.assertAlmostEqual(res["Accuracy"].cpu().item(), self.imagenet_pretrained_accuracies[Models.REGNETY400], delta=0.001)
 
 
     def test_transfer_learning_regnetY400_imagenet(self):
     def test_transfer_learning_regnetY400_imagenet(self):
         trainer = Trainer("imagenet_pretrained_regnetY400_transfer_learning")
         trainer = Trainer("imagenet_pretrained_regnetY400_transfer_learning")
@@ -344,8 +344,8 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer = Trainer("imagenet_pretrained_regnetY200")
         trainer = Trainer("imagenet_pretrained_regnetY200")
 
 
         model = models.get(Models.REGNETY200, arch_params=self.imagenet_pretrained_arch_params["regnet"], **self.imagenet_pretrained_ckpt_params)
         model = models.get(Models.REGNETY200, arch_params=self.imagenet_pretrained_arch_params["regnet"], **self.imagenet_pretrained_ckpt_params)
-        res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)[0].cpu().item()
-        self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies[Models.REGNETY200], delta=0.001)
+        res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)
+        self.assertAlmostEqual(res["Accuracy"].cpu().item(), self.imagenet_pretrained_accuracies[Models.REGNETY200], delta=0.001)
 
 
     def test_transfer_learning_regnetY200_imagenet(self):
     def test_transfer_learning_regnetY200_imagenet(self):
         trainer = Trainer("imagenet_pretrained_regnetY200_transfer_learning")
         trainer = Trainer("imagenet_pretrained_regnetY200_transfer_learning")
@@ -361,8 +361,8 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer = Trainer("imagenet_pretrained_repvgg_a0")
         trainer = Trainer("imagenet_pretrained_repvgg_a0")
 
 
         model = models.get(Models.REPVGG_A0, arch_params=self.imagenet_pretrained_arch_params[Models.REPVGG_A0], **self.imagenet_pretrained_ckpt_params)
         model = models.get(Models.REPVGG_A0, arch_params=self.imagenet_pretrained_arch_params[Models.REPVGG_A0], **self.imagenet_pretrained_ckpt_params)
-        res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)[0].cpu().item()
-        self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies[Models.REPVGG_A0], delta=0.001)
+        res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)
+        self.assertAlmostEqual(res["Accuracy"].cpu().item(), self.imagenet_pretrained_accuracies[Models.REPVGG_A0], delta=0.001)
 
 
     def test_transfer_learning_repvgg_a0_imagenet(self):
     def test_transfer_learning_repvgg_a0_imagenet(self):
         trainer = Trainer("imagenet_pretrained_repvgg_a0_transfer_learning")
         trainer = Trainer("imagenet_pretrained_repvgg_a0_transfer_learning")
@@ -379,14 +379,10 @@ class PretrainedModelsTest(unittest.TestCase):
     def test_pretrained_regseg48_cityscapes(self):
     def test_pretrained_regseg48_cityscapes(self):
         trainer = Trainer("cityscapes_pretrained_regseg48")
         trainer = Trainer("cityscapes_pretrained_regseg48")
         model = models.get(Models.REGSEG48, arch_params=self.cityscapes_pretrained_arch_params[Models.REGSEG48], **self.cityscapes_pretrained_ckpt_params)
         model = models.get(Models.REGSEG48, arch_params=self.cityscapes_pretrained_arch_params[Models.REGSEG48], **self.cityscapes_pretrained_ckpt_params)
-        res = (
-            trainer.test(
-                model=model, test_loader=self.cityscapes_dataset, test_metrics_list=[IoU(num_classes=20, ignore_index=19)], metrics_progress_verbose=True
-            )[0]
-            .cpu()
-            .item()
+        res = trainer.test(
+            model=model, test_loader=self.cityscapes_dataset, test_metrics_list=[IoU(num_classes=20, ignore_index=19)], metrics_progress_verbose=True
         )
         )
-        self.assertAlmostEqual(res, self.cityscapes_pretrained_mious[Models.REGSEG48], delta=0.001)
+        self.assertAlmostEqual(res["IoU"].cpu().item(), self.cityscapes_pretrained_mious[Models.REGSEG48], delta=0.001)
 
 
     def test_transfer_learning_regseg48_cityscapes(self):
     def test_transfer_learning_regseg48_cityscapes(self):
         trainer = Trainer("regseg48_cityscapes_transfer_learning")
         trainer = Trainer("regseg48_cityscapes_transfer_learning")
@@ -401,40 +397,28 @@ class PretrainedModelsTest(unittest.TestCase):
     def test_pretrained_ddrnet23_cityscapes(self):
     def test_pretrained_ddrnet23_cityscapes(self):
         trainer = Trainer("cityscapes_pretrained_ddrnet23")
         trainer = Trainer("cityscapes_pretrained_ddrnet23")
         model = models.get(Models.DDRNET_23, arch_params=self.cityscapes_pretrained_arch_params[Models.DDRNET_23], **self.cityscapes_pretrained_ckpt_params)
         model = models.get(Models.DDRNET_23, arch_params=self.cityscapes_pretrained_arch_params[Models.DDRNET_23], **self.cityscapes_pretrained_ckpt_params)
-        res = (
-            trainer.test(
-                model=model, test_loader=self.cityscapes_dataset, test_metrics_list=[IoU(num_classes=20, ignore_index=19)], metrics_progress_verbose=True
-            )[0]
-            .cpu()
-            .item()
+        res = trainer.test(
+            model=model, test_loader=self.cityscapes_dataset, test_metrics_list=[IoU(num_classes=20, ignore_index=19)], metrics_progress_verbose=True
         )
         )
-        self.assertAlmostEqual(res, self.cityscapes_pretrained_mious[Models.DDRNET_23], delta=0.001)
+        self.assertAlmostEqual(res["IoU"].cpu().item(), self.cityscapes_pretrained_mious[Models.DDRNET_23], delta=0.001)
 
 
     def test_pretrained_ddrnet39_cityscapes(self):
     def test_pretrained_ddrnet39_cityscapes(self):
         trainer = Trainer("cityscapes_pretrained_ddrnet39")
         trainer = Trainer("cityscapes_pretrained_ddrnet39")
         model = models.get(Models.DDRNET_39, arch_params=self.cityscapes_pretrained_arch_params[Models.DDRNET_23], **self.cityscapes_pretrained_ckpt_params)
         model = models.get(Models.DDRNET_39, arch_params=self.cityscapes_pretrained_arch_params[Models.DDRNET_23], **self.cityscapes_pretrained_ckpt_params)
-        res = (
-            trainer.test(
-                model=model, test_loader=self.cityscapes_dataset, test_metrics_list=[IoU(num_classes=20, ignore_index=19)], metrics_progress_verbose=True
-            )[0]
-            .cpu()
-            .item()
+        res = trainer.test(
+            model=model, test_loader=self.cityscapes_dataset, test_metrics_list=[IoU(num_classes=20, ignore_index=19)], metrics_progress_verbose=True
         )
         )
-        self.assertAlmostEqual(res, self.cityscapes_pretrained_mious[Models.DDRNET_39], delta=0.001)
+        self.assertAlmostEqual(res["IoU"].cpu().item(), self.cityscapes_pretrained_mious[Models.DDRNET_39], delta=0.001)
 
 
     def test_pretrained_ddrnet23_slim_cityscapes(self):
     def test_pretrained_ddrnet23_slim_cityscapes(self):
         trainer = Trainer("cityscapes_pretrained_ddrnet23_slim")
         trainer = Trainer("cityscapes_pretrained_ddrnet23_slim")
         model = models.get(
         model = models.get(
             Models.DDRNET_23_SLIM, arch_params=self.cityscapes_pretrained_arch_params[Models.DDRNET_23], **self.cityscapes_pretrained_ckpt_params
             Models.DDRNET_23_SLIM, arch_params=self.cityscapes_pretrained_arch_params[Models.DDRNET_23], **self.cityscapes_pretrained_ckpt_params
         )
         )
-        res = (
-            trainer.test(
-                model=model, test_loader=self.cityscapes_dataset, test_metrics_list=[IoU(num_classes=20, ignore_index=19)], metrics_progress_verbose=True
-            )[0]
-            .cpu()
-            .item()
+        res = trainer.test(
+            model=model, test_loader=self.cityscapes_dataset, test_metrics_list=[IoU(num_classes=20, ignore_index=19)], metrics_progress_verbose=True
         )
         )
-        self.assertAlmostEqual(res, self.cityscapes_pretrained_mious[Models.DDRNET_23_SLIM], delta=0.001)
+        self.assertAlmostEqual(res["IoU"].cpu().item(), self.cityscapes_pretrained_mious[Models.DDRNET_23_SLIM], delta=0.001)
 
 
     def test_transfer_learning_ddrnet23_cityscapes(self):
     def test_transfer_learning_ddrnet23_cityscapes(self):
         trainer = Trainer("cityscapes_pretrained_ddrnet23_transfer_learning")
         trainer = Trainer("cityscapes_pretrained_ddrnet23_transfer_learning")
@@ -465,8 +449,8 @@ class PretrainedModelsTest(unittest.TestCase):
             arch_params=self.coco_segmentation_subclass_pretrained_arch_params["shelfnet34_lw"],
             arch_params=self.coco_segmentation_subclass_pretrained_arch_params["shelfnet34_lw"],
             **self.coco_segmentation_subclass_pretrained_ckpt_params,
             **self.coco_segmentation_subclass_pretrained_ckpt_params,
         )
         )
-        res = trainer.test(model=model, test_loader=self.coco_segmentation_dataset, test_metrics_list=[IoU(21)], metrics_progress_verbose=True)[0].cpu().item()
-        self.assertAlmostEqual(res, self.coco_segmentation_subclass_pretrained_mious["shelfnet34_lw"], delta=0.001)
+        res = trainer.test(model=model, test_loader=self.coco_segmentation_dataset, test_metrics_list=[IoU(21)], metrics_progress_verbose=True)
+        self.assertAlmostEqual(res["IoU"].cpu().item(), self.cityscapes_pretrained_mious["shelfnet34_lw"], delta=0.001)
 
 
     def test_pretrained_efficientnet_b0_imagenet(self):
     def test_pretrained_efficientnet_b0_imagenet(self):
         trainer = Trainer("imagenet_pretrained_efficientnet_b0")
         trainer = Trainer("imagenet_pretrained_efficientnet_b0")
@@ -474,8 +458,8 @@ class PretrainedModelsTest(unittest.TestCase):
         model = models.get(
         model = models.get(
             Models.EFFICIENTNET_B0, arch_params=self.imagenet_pretrained_arch_params[Models.EFFICIENTNET_B0], **self.imagenet_pretrained_ckpt_params
             Models.EFFICIENTNET_B0, arch_params=self.imagenet_pretrained_arch_params[Models.EFFICIENTNET_B0], **self.imagenet_pretrained_ckpt_params
         )
         )
-        res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)[0].cpu().item()
-        self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies[Models.EFFICIENTNET_B0], delta=0.001)
+        res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)
+        self.assertAlmostEqual(res["Accuracy"].cpu().item(), self.imagenet_pretrained_accuracies[Models.EFFICIENTNET_B0], delta=0.001)
 
 
     def test_transfer_learning_efficientnet_b0_imagenet(self):
     def test_transfer_learning_efficientnet_b0_imagenet(self):
         trainer = Trainer("imagenet_pretrained_efficientnet_b0_transfer_learning")
         trainer = Trainer("imagenet_pretrained_efficientnet_b0_transfer_learning")
@@ -504,8 +488,8 @@ class PretrainedModelsTest(unittest.TestCase):
             test_loader=self.coco_dataset["ssd_mobilenet"],
             test_loader=self.coco_dataset["ssd_mobilenet"],
             test_metrics_list=[DetectionMetrics(post_prediction_callback=ssd_post_prediction_callback, num_cls=80)],
             test_metrics_list=[DetectionMetrics(post_prediction_callback=ssd_post_prediction_callback, num_cls=80)],
             metrics_progress_verbose=True,
             metrics_progress_verbose=True,
-        )[2]
-        self.assertAlmostEqual(res, self.coco_pretrained_maps[Models.SSD_LITE_MOBILENET_V2], delta=0.001)
+        )
+        self.assertAlmostEqual(res["mAP@0.50:0.95"].cpu().item(), self.coco_pretrained_maps[Models.SSD_LITE_MOBILENET_V2], delta=0.001)
 
 
     def test_transfer_learning_ssd_lite_mobilenet_v2_coco(self):
     def test_transfer_learning_ssd_lite_mobilenet_v2_coco(self):
         trainer = Trainer("coco_ssd_lite_mobilenet_v2_transfer_learning")
         trainer = Trainer("coco_ssd_lite_mobilenet_v2_transfer_learning")
@@ -528,8 +512,8 @@ class PretrainedModelsTest(unittest.TestCase):
             test_loader=self.coco_dataset["ssd_mobilenet"],
             test_loader=self.coco_dataset["ssd_mobilenet"],
             test_metrics_list=[DetectionMetrics(post_prediction_callback=ssd_post_prediction_callback, num_cls=80)],
             test_metrics_list=[DetectionMetrics(post_prediction_callback=ssd_post_prediction_callback, num_cls=80)],
             metrics_progress_verbose=True,
             metrics_progress_verbose=True,
-        )[2]
-        self.assertAlmostEqual(res, self.coco_pretrained_maps[Models.SSD_MOBILENET_V1], delta=0.001)
+        )
+        self.assertAlmostEqual(res["mAP@0.50:0.95"].cpu().item(), self.coco_pretrained_maps[Models.SSD_MOBILENET_V1], delta=0.001)
 
 
     def test_pretrained_yolox_s_coco(self):
     def test_pretrained_yolox_s_coco(self):
         trainer = Trainer(Models.YOLOX_S)
         trainer = Trainer(Models.YOLOX_S)
@@ -539,8 +523,8 @@ class PretrainedModelsTest(unittest.TestCase):
             model=model,
             model=model,
             test_loader=self.coco_dataset["yolox"],
             test_loader=self.coco_dataset["yolox"],
             test_metrics_list=[DetectionMetrics(post_prediction_callback=YoloPostPredictionCallback(), num_cls=80, normalize_targets=True)],
             test_metrics_list=[DetectionMetrics(post_prediction_callback=YoloPostPredictionCallback(), num_cls=80, normalize_targets=True)],
-        )[2]
-        self.assertAlmostEqual(res, self.coco_pretrained_maps[Models.YOLOX_S], delta=0.001)
+        )
+        self.assertAlmostEqual(res["mAP@0.50:0.95"].cpu().item(), self.coco_pretrained_maps[Models.YOLOX_S], delta=0.001)
 
 
     def test_pretrained_yolox_m_coco(self):
     def test_pretrained_yolox_m_coco(self):
         trainer = Trainer(Models.YOLOX_M)
         trainer = Trainer(Models.YOLOX_M)
@@ -549,8 +533,8 @@ class PretrainedModelsTest(unittest.TestCase):
             model=model,
             model=model,
             test_loader=self.coco_dataset["yolox"],
             test_loader=self.coco_dataset["yolox"],
             test_metrics_list=[DetectionMetrics(post_prediction_callback=YoloPostPredictionCallback(), num_cls=80, normalize_targets=True)],
             test_metrics_list=[DetectionMetrics(post_prediction_callback=YoloPostPredictionCallback(), num_cls=80, normalize_targets=True)],
-        )[2]
-        self.assertAlmostEqual(res, self.coco_pretrained_maps[Models.YOLOX_M], delta=0.001)
+        )
+        self.assertAlmostEqual(res["mAP@0.50:0.95"].cpu().item(), self.coco_pretrained_maps[Models.YOLOX_M], delta=0.001)
 
 
     def test_pretrained_yolox_l_coco(self):
     def test_pretrained_yolox_l_coco(self):
         trainer = Trainer(Models.YOLOX_L)
         trainer = Trainer(Models.YOLOX_L)
@@ -559,8 +543,8 @@ class PretrainedModelsTest(unittest.TestCase):
             model=model,
             model=model,
             test_loader=self.coco_dataset["yolox"],
             test_loader=self.coco_dataset["yolox"],
             test_metrics_list=[DetectionMetrics(post_prediction_callback=YoloPostPredictionCallback(), num_cls=80, normalize_targets=True)],
             test_metrics_list=[DetectionMetrics(post_prediction_callback=YoloPostPredictionCallback(), num_cls=80, normalize_targets=True)],
-        )[2]
-        self.assertAlmostEqual(res, self.coco_pretrained_maps[Models.YOLOX_L], delta=0.001)
+        )
+        self.assertAlmostEqual(res["mAP@0.50:0.95"].cpu().item(), self.coco_pretrained_maps[Models.YOLOX_L], delta=0.001)
 
 
     def test_pretrained_yolox_n_coco(self):
     def test_pretrained_yolox_n_coco(self):
         trainer = Trainer(Models.YOLOX_N)
         trainer = Trainer(Models.YOLOX_N)
@@ -570,8 +554,8 @@ class PretrainedModelsTest(unittest.TestCase):
             model=model,
             model=model,
             test_loader=self.coco_dataset["yolox"],
             test_loader=self.coco_dataset["yolox"],
             test_metrics_list=[DetectionMetrics(post_prediction_callback=YoloPostPredictionCallback(), num_cls=80, normalize_targets=True)],
             test_metrics_list=[DetectionMetrics(post_prediction_callback=YoloPostPredictionCallback(), num_cls=80, normalize_targets=True)],
-        )[2]
-        self.assertAlmostEqual(res, self.coco_pretrained_maps[Models.YOLOX_N], delta=0.001)
+        )
+        self.assertAlmostEqual(res["mAP@0.50:0.95"].cpu().item(), self.coco_pretrained_maps[Models.YOLOX_N], delta=0.001)
 
 
     def test_pretrained_yolox_t_coco(self):
     def test_pretrained_yolox_t_coco(self):
         trainer = Trainer(Models.YOLOX_T)
         trainer = Trainer(Models.YOLOX_T)
@@ -580,8 +564,8 @@ class PretrainedModelsTest(unittest.TestCase):
             model=model,
             model=model,
             test_loader=self.coco_dataset["yolox"],
             test_loader=self.coco_dataset["yolox"],
             test_metrics_list=[DetectionMetrics(post_prediction_callback=YoloPostPredictionCallback(), num_cls=80, normalize_targets=True)],
             test_metrics_list=[DetectionMetrics(post_prediction_callback=YoloPostPredictionCallback(), num_cls=80, normalize_targets=True)],
-        )[2]
-        self.assertAlmostEqual(res, self.coco_pretrained_maps[Models.YOLOX_T], delta=0.001)
+        )
+        self.assertAlmostEqual(res["mAP@0.50:0.95"].cpu().item(), self.coco_pretrained_maps[Models.YOLOX_T], delta=0.001)
 
 
     def test_pretrained_ppyoloe_s_coco(self):
     def test_pretrained_ppyoloe_s_coco(self):
         trainer = Trainer(Models.PP_YOLOE_S)
         trainer = Trainer(Models.PP_YOLOE_S)
@@ -599,8 +583,8 @@ class PretrainedModelsTest(unittest.TestCase):
                     post_prediction_callback=PPYoloEPostPredictionCallback(score_threshold=0.01, nms_top_k=1000, max_predictions=300, nms_threshold=0.7),
                     post_prediction_callback=PPYoloEPostPredictionCallback(score_threshold=0.01, nms_top_k=1000, max_predictions=300, nms_threshold=0.7),
                 )
                 )
             ],
             ],
-        )[2]
-        self.assertAlmostEqual(res, self.coco_pretrained_maps[Models.PP_YOLOE_S], delta=0.001)
+        )
+        self.assertAlmostEqual(res["mAP@0.50:0.95"].cpu().item(), self.coco_pretrained_maps[Models.PP_YOLOE_S], delta=0.001)
 
 
     def test_pretrained_ppyoloe_m_coco(self):
     def test_pretrained_ppyoloe_m_coco(self):
         trainer = Trainer(Models.PP_YOLOE_M)
         trainer = Trainer(Models.PP_YOLOE_M)
@@ -618,8 +602,8 @@ class PretrainedModelsTest(unittest.TestCase):
                     post_prediction_callback=PPYoloEPostPredictionCallback(score_threshold=0.01, nms_top_k=1000, max_predictions=300, nms_threshold=0.7),
                     post_prediction_callback=PPYoloEPostPredictionCallback(score_threshold=0.01, nms_top_k=1000, max_predictions=300, nms_threshold=0.7),
                 )
                 )
             ],
             ],
-        )[2]
-        self.assertAlmostEqual(res, self.coco_pretrained_maps[Models.PP_YOLOE_M], delta=0.001)
+        )
+        self.assertAlmostEqual(res["mAP@0.50:0.95"].cpu().item(), self.coco_pretrained_maps[Models.PP_YOLOE_M], delta=0.001)
 
 
     def test_pretrained_ppyoloe_l_coco(self):
     def test_pretrained_ppyoloe_l_coco(self):
         trainer = Trainer(Models.PP_YOLOE_L)
         trainer = Trainer(Models.PP_YOLOE_L)
@@ -637,8 +621,8 @@ class PretrainedModelsTest(unittest.TestCase):
                     post_prediction_callback=PPYoloEPostPredictionCallback(score_threshold=0.01, nms_top_k=1000, max_predictions=300, nms_threshold=0.7),
                     post_prediction_callback=PPYoloEPostPredictionCallback(score_threshold=0.01, nms_top_k=1000, max_predictions=300, nms_threshold=0.7),
                 )
                 )
             ],
             ],
-        )[2]
-        self.assertAlmostEqual(res, self.coco_pretrained_maps[Models.PP_YOLOE_L], delta=0.001)
+        )
+        self.assertAlmostEqual(res["mAP@0.50:0.95"].cpu().item(), self.coco_pretrained_maps[Models.PP_YOLOE_L], delta=0.001)
 
 
     def test_pretrained_ppyoloe_x_coco(self):
     def test_pretrained_ppyoloe_x_coco(self):
         trainer = Trainer(Models.PP_YOLOE_X)
         trainer = Trainer(Models.PP_YOLOE_X)
@@ -656,8 +640,8 @@ class PretrainedModelsTest(unittest.TestCase):
                     post_prediction_callback=PPYoloEPostPredictionCallback(score_threshold=0.01, nms_top_k=1000, max_predictions=300, nms_threshold=0.7),
                     post_prediction_callback=PPYoloEPostPredictionCallback(score_threshold=0.01, nms_top_k=1000, max_predictions=300, nms_threshold=0.7),
                 )
                 )
             ],
             ],
-        )[2]
-        self.assertAlmostEqual(res, self.coco_pretrained_maps[Models.PP_YOLOE_X], delta=0.001)
+        )
+        self.assertAlmostEqual(res["mAP@0.50:0.95"].cpu().item(), self.coco_pretrained_maps[Models.PP_YOLOE_X], delta=0.001)
 
 
     def test_transfer_learning_yolox_n_coco(self):
     def test_transfer_learning_yolox_n_coco(self):
         trainer = Trainer("test_transfer_learning_yolox_n_coco")
         trainer = Trainer("test_transfer_learning_yolox_n_coco")
@@ -686,8 +670,8 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer = Trainer("imagenet_mobilenet_v3_large")
         trainer = Trainer("imagenet_mobilenet_v3_large")
 
 
         model = models.get(Models.MOBILENET_V3_LARGE, arch_params=self.imagenet_pretrained_arch_params["mobilenet"], **self.imagenet_pretrained_ckpt_params)
         model = models.get(Models.MOBILENET_V3_LARGE, arch_params=self.imagenet_pretrained_arch_params["mobilenet"], **self.imagenet_pretrained_ckpt_params)
-        res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)[0].cpu().item()
-        self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies[Models.MOBILENET_V3_LARGE], delta=0.001)
+        res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)
+        self.assertAlmostEqual(res["Accuracy"].cpu().item(), self.imagenet_pretrained_accuracies[Models.MOBILENET_V3_LARGE], delta=0.001)
 
 
     def test_transfer_learning_mobilenet_v3_small_imagenet(self):
     def test_transfer_learning_mobilenet_v3_small_imagenet(self):
         trainer = Trainer("imagenet_pretrained_mobilenet_v3_small_transfer_learning")
         trainer = Trainer("imagenet_pretrained_mobilenet_v3_small_transfer_learning")
@@ -706,8 +690,8 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer = Trainer("imagenet_mobilenet_v3_small")
         trainer = Trainer("imagenet_mobilenet_v3_small")
 
 
         model = models.get(Models.MOBILENET_V3_SMALL, arch_params=self.imagenet_pretrained_arch_params["mobilenet"], **self.imagenet_pretrained_ckpt_params)
         model = models.get(Models.MOBILENET_V3_SMALL, arch_params=self.imagenet_pretrained_arch_params["mobilenet"], **self.imagenet_pretrained_ckpt_params)
-        res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)[0].cpu().item()
-        self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies[Models.MOBILENET_V3_SMALL], delta=0.001)
+        res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)
+        self.assertAlmostEqual(res["Accuracy"].cpu().item(), self.imagenet_pretrained_accuracies[Models.MOBILENET_V3_SMALL], delta=0.001)
 
 
     def test_transfer_learning_mobilenet_v2_imagenet(self):
     def test_transfer_learning_mobilenet_v2_imagenet(self):
         trainer = Trainer("imagenet_pretrained_mobilenet_v2_transfer_learning")
         trainer = Trainer("imagenet_pretrained_mobilenet_v2_transfer_learning")
@@ -726,24 +710,20 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer = Trainer("imagenet_mobilenet_v2")
         trainer = Trainer("imagenet_mobilenet_v2")
 
 
         model = models.get(Models.MOBILENET_V2, arch_params=self.imagenet_pretrained_arch_params["mobilenet"], **self.imagenet_pretrained_ckpt_params)
         model = models.get(Models.MOBILENET_V2, arch_params=self.imagenet_pretrained_arch_params["mobilenet"], **self.imagenet_pretrained_ckpt_params)
-        res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)[0].cpu().item()
-        self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies[Models.MOBILENET_V2], delta=0.001)
+        res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)
+        self.assertAlmostEqual(res["Accuracy"].cpu().item(), self.imagenet_pretrained_accuracies[Models.MOBILENET_V2], delta=0.001)
 
 
     def test_pretrained_stdc1_seg50_cityscapes(self):
     def test_pretrained_stdc1_seg50_cityscapes(self):
         trainer = Trainer("cityscapes_pretrained_stdc1_seg50")
         trainer = Trainer("cityscapes_pretrained_stdc1_seg50")
 
 
         model = models.get(Models.STDC1_SEG50, arch_params=self.cityscapes_pretrained_arch_params["stdc"], **self.cityscapes_pretrained_ckpt_params)
         model = models.get(Models.STDC1_SEG50, arch_params=self.cityscapes_pretrained_arch_params["stdc"], **self.cityscapes_pretrained_ckpt_params)
-        res = (
-            trainer.test(
-                model=model,
-                test_loader=self.cityscapes_dataset_rescaled50,
-                test_metrics_list=[IoU(num_classes=20, ignore_index=19)],
-                metrics_progress_verbose=True,
-            )[0]
-            .cpu()
-            .item()
-        )
-        self.assertAlmostEqual(res, self.cityscapes_pretrained_mious[Models.STDC1_SEG50], delta=0.001)
+        res = trainer.test(
+            model=model,
+            test_loader=self.cityscapes_dataset_rescaled50,
+            test_metrics_list=[IoU(num_classes=20, ignore_index=19)],
+            metrics_progress_verbose=True,
+        )
+        self.assertAlmostEqual(res["IoU"].cpu().item(), self.cityscapes_pretrained_mious[Models.STDC1_SEG50], delta=0.001)
 
 
     def test_transfer_learning_stdc1_seg50_cityscapes(self):
     def test_transfer_learning_stdc1_seg50_cityscapes(self):
         trainer = Trainer("cityscapes_pretrained_stdc1_seg50_transfer_learning")
         trainer = Trainer("cityscapes_pretrained_stdc1_seg50_transfer_learning")
@@ -760,17 +740,13 @@ class PretrainedModelsTest(unittest.TestCase):
     def test_pretrained_stdc1_seg75_cityscapes(self):
     def test_pretrained_stdc1_seg75_cityscapes(self):
         trainer = Trainer("cityscapes_pretrained_stdc1_seg75")
         trainer = Trainer("cityscapes_pretrained_stdc1_seg75")
         model = models.get(Models.STDC1_SEG75, arch_params=self.cityscapes_pretrained_arch_params["stdc"], **self.cityscapes_pretrained_ckpt_params)
         model = models.get(Models.STDC1_SEG75, arch_params=self.cityscapes_pretrained_arch_params["stdc"], **self.cityscapes_pretrained_ckpt_params)
-        res = (
-            trainer.test(
-                model=model,
-                test_loader=self.cityscapes_dataset_rescaled75,
-                test_metrics_list=[IoU(num_classes=20, ignore_index=19)],
-                metrics_progress_verbose=True,
-            )[0]
-            .cpu()
-            .item()
-        )
-        self.assertAlmostEqual(res, self.cityscapes_pretrained_mious[Models.STDC1_SEG75], delta=0.001)
+        res = trainer.test(
+            model=model,
+            test_loader=self.cityscapes_dataset_rescaled75,
+            test_metrics_list=[IoU(num_classes=20, ignore_index=19)],
+            metrics_progress_verbose=True,
+        )
+        self.assertAlmostEqual(res["IoU"].cpu().item(), self.cityscapes_pretrained_mious[Models.STDC1_SEG75], delta=0.001)
 
 
     def test_transfer_learning_stdc1_seg75_cityscapes(self):
     def test_transfer_learning_stdc1_seg75_cityscapes(self):
         trainer = Trainer("cityscapes_pretrained_stdc1_seg75_transfer_learning")
         trainer = Trainer("cityscapes_pretrained_stdc1_seg75_transfer_learning")
@@ -787,17 +763,10 @@ class PretrainedModelsTest(unittest.TestCase):
     def test_pretrained_stdc2_seg50_cityscapes(self):
     def test_pretrained_stdc2_seg50_cityscapes(self):
         trainer = Trainer("cityscapes_pretrained_stdc2_seg50")
         trainer = Trainer("cityscapes_pretrained_stdc2_seg50")
         model = models.get(Models.STDC2_SEG50, arch_params=self.cityscapes_pretrained_arch_params["stdc"], **self.cityscapes_pretrained_ckpt_params)
         model = models.get(Models.STDC2_SEG50, arch_params=self.cityscapes_pretrained_arch_params["stdc"], **self.cityscapes_pretrained_ckpt_params)
-        res = (
-            trainer.test(
-                model=model,
-                test_loader=self.cityscapes_dataset_rescaled50,
-                test_metrics_list=[IoU(num_classes=20, ignore_index=19)],
-                metrics_progress_verbose=True,
-            )[0]
-            .cpu()
-            .item()
-        )
-        self.assertAlmostEqual(res, self.cityscapes_pretrained_mious[Models.STDC2_SEG50], delta=0.001)
+        res = trainer.test(
+            model=model, test_loader=self.cityscapes_dataset_rescaled50, test_metrics_list=[IoU(num_classes=20, ignore_index=19)], metrics_progress_verbose=True
+        )
+        self.assertAlmostEqual(res["IoU"].cpu().item(), self.cityscapes_pretrained_mious[Models.STDC2_SEG50], delta=0.001)
 
 
     def test_transfer_learning_stdc2_seg50_cityscapes(self):
     def test_transfer_learning_stdc2_seg50_cityscapes(self):
         trainer = Trainer("cityscapes_pretrained_stdc2_seg50_transfer_learning")
         trainer = Trainer("cityscapes_pretrained_stdc2_seg50_transfer_learning")
@@ -814,17 +783,13 @@ class PretrainedModelsTest(unittest.TestCase):
     def test_pretrained_stdc2_seg75_cityscapes(self):
     def test_pretrained_stdc2_seg75_cityscapes(self):
         trainer = Trainer("cityscapes_pretrained_stdc2_seg75")
         trainer = Trainer("cityscapes_pretrained_stdc2_seg75")
         model = models.get(Models.STDC2_SEG75, arch_params=self.cityscapes_pretrained_arch_params["stdc"], **self.cityscapes_pretrained_ckpt_params)
         model = models.get(Models.STDC2_SEG75, arch_params=self.cityscapes_pretrained_arch_params["stdc"], **self.cityscapes_pretrained_ckpt_params)
-        res = (
-            trainer.test(
-                model=model,
-                test_loader=self.cityscapes_dataset_rescaled75,
-                test_metrics_list=[IoU(num_classes=20, ignore_index=19)],
-                metrics_progress_verbose=True,
-            )[0]
-            .cpu()
-            .item()
-        )
-        self.assertAlmostEqual(res, self.cityscapes_pretrained_mious[Models.STDC2_SEG75], delta=0.001)
+        res = trainer.test(
+            model=model,
+            test_loader=self.cityscapes_dataset_rescaled75,
+            test_metrics_list=[IoU(num_classes=20, ignore_index=19)],
+            metrics_progress_verbose=True,
+        )
+        self.assertAlmostEqual(res["IoU"].cpu().item(), self.cityscapes_pretrained_mious[Models.STDC2_SEG75], delta=0.001)
 
 
     def test_transfer_learning_stdc2_seg75_cityscapes(self):
     def test_transfer_learning_stdc2_seg75_cityscapes(self):
         trainer = Trainer("cityscapes_pretrained_stdc2_seg75_transfer_learning")
         trainer = Trainer("cityscapes_pretrained_stdc2_seg75_transfer_learning")
@@ -867,34 +832,22 @@ class PretrainedModelsTest(unittest.TestCase):
     def test_pretrained_vit_base_imagenet(self):
     def test_pretrained_vit_base_imagenet(self):
         trainer = Trainer("imagenet_pretrained_vit_base")
         trainer = Trainer("imagenet_pretrained_vit_base")
         model = models.get(Models.VIT_BASE, arch_params=self.imagenet_pretrained_arch_params[Models.VIT_BASE], **self.imagenet_pretrained_ckpt_params)
         model = models.get(Models.VIT_BASE, arch_params=self.imagenet_pretrained_arch_params[Models.VIT_BASE], **self.imagenet_pretrained_ckpt_params)
-        res = (
-            trainer.test(model=model, test_loader=self.imagenet_dataset_05_mean_std, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)[0]
-            .cpu()
-            .item()
-        )
-        self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies[Models.VIT_BASE], delta=0.001)
+        res = trainer.test(model=model, test_loader=self.imagenet_dataset_05_mean_std, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)
+        self.assertAlmostEqual(res["Accuracy"].cpu().item(), self.imagenet_pretrained_accuracies[Models.VIT_BASE], delta=0.001)
 
 
     def test_pretrained_vit_large_imagenet(self):
     def test_pretrained_vit_large_imagenet(self):
         trainer = Trainer("imagenet_pretrained_vit_large")
         trainer = Trainer("imagenet_pretrained_vit_large")
         model = models.get(Models.VIT_LARGE, arch_params=self.imagenet_pretrained_arch_params[Models.VIT_BASE], **self.imagenet_pretrained_ckpt_params)
         model = models.get(Models.VIT_LARGE, arch_params=self.imagenet_pretrained_arch_params[Models.VIT_BASE], **self.imagenet_pretrained_ckpt_params)
-        res = (
-            trainer.test(model=model, test_loader=self.imagenet_dataset_05_mean_std, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)[0]
-            .cpu()
-            .item()
-        )
-        self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies[Models.VIT_LARGE], delta=0.001)
+        res = trainer.test(model=model, test_loader=self.imagenet_dataset_05_mean_std, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)
+        self.assertAlmostEqual(res["Accuracy"].cpu().item(), self.imagenet_pretrained_accuracies[Models.VIT_LARGE], delta=0.001)
 
 
     def test_pretrained_beit_base_imagenet(self):
     def test_pretrained_beit_base_imagenet(self):
         trainer = Trainer("imagenet_pretrained_beit_base")
         trainer = Trainer("imagenet_pretrained_beit_base")
         model = models.get(
         model = models.get(
             Models.BEIT_BASE_PATCH16_224, arch_params=self.imagenet_pretrained_arch_params[Models.VIT_BASE], **self.imagenet_pretrained_ckpt_params
             Models.BEIT_BASE_PATCH16_224, arch_params=self.imagenet_pretrained_arch_params[Models.VIT_BASE], **self.imagenet_pretrained_ckpt_params
         )
         )
-        res = (
-            trainer.test(model=model, test_loader=self.imagenet_dataset_05_mean_std, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)[0]
-            .cpu()
-            .item()
-        )
-        self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies[Models.BEIT_BASE_PATCH16_224], delta=0.001)
+        res = trainer.test(model=model, test_loader=self.imagenet_dataset_05_mean_std, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)
+        self.assertAlmostEqual(res["Accuracy"].cpu().item(), self.imagenet_pretrained_accuracies[Models.BEIT_BASE_PATCH16_224], delta=0.001)
 
 
     def test_transfer_learning_beit_base_imagenet(self):
     def test_transfer_learning_beit_base_imagenet(self):
         trainer = Trainer("test_transfer_learning_beit_base_imagenet")
         trainer = Trainer("test_transfer_learning_beit_base_imagenet")
@@ -916,65 +869,43 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer = Trainer("cityscapes_pretrained_pplite_t_seg50")
         trainer = Trainer("cityscapes_pretrained_pplite_t_seg50")
         model = models.get(Models.PP_LITE_T_SEG50, arch_params=self.cityscapes_pretrained_arch_params["pplite_seg"], **self.cityscapes_pretrained_ckpt_params)
         model = models.get(Models.PP_LITE_T_SEG50, arch_params=self.cityscapes_pretrained_arch_params["pplite_seg"], **self.cityscapes_pretrained_ckpt_params)
 
 
-        res = (
-            trainer.test(
-                model=model,
-                test_loader=self.cityscapes_dataset_rescaled50,
-                test_metrics_list=[IoU(num_classes=20, ignore_index=19)],
-                metrics_progress_verbose=True,
-            )[0]
-            .cpu()
-            .item()
+        res = trainer.test(
+            model=model, test_loader=self.cityscapes_dataset_rescaled50, test_metrics_list=[IoU(num_classes=20, ignore_index=19)], metrics_progress_verbose=True
         )
         )
-        self.assertAlmostEqual(res, self.cityscapes_pretrained_mious[Models.PP_LITE_T_SEG50], delta=0.001)
+        self.assertAlmostEqual(res["IoU"].cpu().item(), self.cityscapes_pretrained_mious[Models.PP_LITE_T_SEG50], delta=0.001)
 
 
     def test_pretrained_pplite_t_seg75_cityscapes(self):
     def test_pretrained_pplite_t_seg75_cityscapes(self):
         trainer = Trainer("cityscapes_pretrained_pplite_t_seg75")
         trainer = Trainer("cityscapes_pretrained_pplite_t_seg75")
         model = models.get(Models.PP_LITE_T_SEG75, arch_params=self.cityscapes_pretrained_arch_params["pplite_seg"], **self.cityscapes_pretrained_ckpt_params)
         model = models.get(Models.PP_LITE_T_SEG75, arch_params=self.cityscapes_pretrained_arch_params["pplite_seg"], **self.cityscapes_pretrained_ckpt_params)
 
 
-        res = (
-            trainer.test(
-                model=model,
-                test_loader=self.cityscapes_dataset_rescaled50,
-                test_metrics_list=[IoU(num_classes=20, ignore_index=19)],
-                metrics_progress_verbose=True,
-            )[0]
-            .cpu()
-            .item()
+        res = trainer.test(
+            model=model, test_loader=self.cityscapes_dataset_rescaled50, test_metrics_list=[IoU(num_classes=20, ignore_index=19)], metrics_progress_verbose=True
         )
         )
-        self.assertAlmostEqual(res, self.cityscapes_pretrained_mious[Models.PP_LITE_T_SEG75], delta=0.001)
+        self.assertAlmostEqual(res["IoU"].cpu().item(), self.cityscapes_pretrained_mious[Models.PP_LITE_T_SEG75], delta=0.001)
 
 
     def test_pretrained_pplite_b_seg50_cityscapes(self):
     def test_pretrained_pplite_b_seg50_cityscapes(self):
         trainer = Trainer("cityscapes_pretrained_pplite_b_seg50")
         trainer = Trainer("cityscapes_pretrained_pplite_b_seg50")
         model = models.get(Models.PP_LITE_B_SEG50, arch_params=self.cityscapes_pretrained_arch_params["pplite_seg"], **self.cityscapes_pretrained_ckpt_params)
         model = models.get(Models.PP_LITE_B_SEG50, arch_params=self.cityscapes_pretrained_arch_params["pplite_seg"], **self.cityscapes_pretrained_ckpt_params)
 
 
-        res = (
-            trainer.test(
-                model=model,
-                test_loader=self.cityscapes_dataset_rescaled50,
-                test_metrics_list=[IoU(num_classes=20, ignore_index=19)],
-                metrics_progress_verbose=True,
-            )[0]
-            .cpu()
-            .item()
+        res = trainer.test(
+            model=model,
+            test_loader=self.cityscapes_dataset_rescaled50,
+            test_metrics_list=[IoU(num_classes=20, ignore_index=19)],
+            metrics_progress_verbose=True,
         )
         )
-        self.assertAlmostEqual(res, self.cityscapes_pretrained_mious[Models.PP_LITE_B_SEG50], delta=0.001)
+        self.assertAlmostEqual(res["IoU"].cpu().item(), self.cityscapes_pretrained_mious[Models.PP_LITE_B_SEG50], delta=0.001)
 
 
     def test_pretrained_pplite_b_seg75_cityscapes(self):
     def test_pretrained_pplite_b_seg75_cityscapes(self):
         trainer = Trainer("cityscapes_pretrained_pplite_b_seg75")
         trainer = Trainer("cityscapes_pretrained_pplite_b_seg75")
         model = models.get(Models.PP_LITE_B_SEG75, arch_params=self.cityscapes_pretrained_arch_params["pplite_seg"], **self.cityscapes_pretrained_ckpt_params)
         model = models.get(Models.PP_LITE_B_SEG75, arch_params=self.cityscapes_pretrained_arch_params["pplite_seg"], **self.cityscapes_pretrained_ckpt_params)
 
 
-        res = (
-            trainer.test(
-                model=model,
-                test_loader=self.cityscapes_dataset_rescaled50,
-                test_metrics_list=[IoU(num_classes=20, ignore_index=19)],
-                metrics_progress_verbose=True,
-            )[0]
-            .cpu()
-            .item()
-        )
-        self.assertAlmostEqual(res, self.cityscapes_pretrained_mious[Models.PP_LITE_B_SEG75], delta=0.001)
+        res = trainer.test(
+            model=model,
+            test_loader=self.cityscapes_dataset_rescaled50,
+            test_metrics_list=[IoU(num_classes=20, ignore_index=19)],
+            metrics_progress_verbose=True,
+        )
+        self.assertAlmostEqual(res["IoU"].cpu().item(), self.cityscapes_pretrained_mious[Models.PP_LITE_B_SEG75], delta=0.001)
 
 
     def tearDown(self) -> None:
     def tearDown(self) -> None:
         if os.path.exists("~/.cache/torch/hub/"):
         if os.path.exists("~/.cache/torch/hub/"):
Discard
@@ -46,7 +46,7 @@ class TestWithoutTrainTest(unittest.TestCase):
     def test_test_without_train(self):
     def test_test_without_train(self):
         trainer, model = self.get_classification_trainer(self.folder_names[0])
         trainer, model = self.get_classification_trainer(self.folder_names[0])
         assert isinstance(
         assert isinstance(
-            trainer.test(model=model, silent_mode=True, test_metrics_list=[Accuracy(), Top5()], test_loader=classification_test_dataloader()), tuple
+            trainer.test(model=model, silent_mode=True, test_metrics_list=[Accuracy(), Top5()], test_loader=classification_test_dataloader()), dict
         )
         )
 
 
         trainer, model = self.get_detection_trainer(self.folder_names[1])
         trainer, model = self.get_detection_trainer(self.folder_names[1])
@@ -54,12 +54,12 @@ class TestWithoutTrainTest(unittest.TestCase):
         test_metrics = [DetectionMetrics(post_prediction_callback=YoloPostPredictionCallback(), num_cls=5)]
         test_metrics = [DetectionMetrics(post_prediction_callback=YoloPostPredictionCallback(), num_cls=5)]
 
 
         assert isinstance(
         assert isinstance(
-            trainer.test(model=model, silent_mode=True, test_metrics_list=test_metrics, test_loader=detection_test_dataloader(image_size=320)), tuple
+            trainer.test(model=model, silent_mode=True, test_metrics_list=test_metrics, test_loader=detection_test_dataloader(image_size=320)), dict
         )
         )
 
 
         trainer, model = self.get_segmentation_trainer(self.folder_names[2])
         trainer, model = self.get_segmentation_trainer(self.folder_names[2])
         assert isinstance(
         assert isinstance(
-            trainer.test(model=model, silent_mode=True, test_metrics_list=[IoU(21), PixelAccuracy()], test_loader=segmentation_test_dataloader()), tuple
+            trainer.test(model=model, silent_mode=True, test_metrics_list=[IoU(21), PixelAccuracy()], test_loader=segmentation_test_dataloader()), dict
         )
         )
 
 
 
 
Discard