Deci-AI
/
super-gradients
connected to https://github.com/Deci-AI/super-gradients.git


  
1

	
2

	
3

	
4

	
5

	
6

	
7

	
8

	
9

	
10

	
11

	
12

	
13

	
14

	
15

	
16

	
17

	
18

	
19

	
20

	
21

	
22

	
23

	
24

	
25

	
26

	
27

	
28

	
29

	
30

	
31

	
32

	
33

	
34

	
35

	
36

	
37

	
38

	
39

	
40

	
41

	
42

	
43

	
44

	
45

	
46

	
47

	
48

	
49

	
50

	
51

	
52

	
53

	
54

	
55

	
56

	
57

	
58

	
59

	
60

	
61

	
62

	
63

	
64

	
65

	
66

	
67

	
68

	
69

	
70

	
71

	
72

	
73

	
74

	
75

	
76

	
77

	
78

	
79

	
80

	
81

	
82

	
83

	
84

	
85

	
86

	
87

	
88

	
89

	
90

	
91

	
92

	
93

	
94

	
            import shutil
import unittest
import os
from super_gradients import Trainer, \
    ClassificationTestDatasetInterface, \
    SegmentationTestDatasetInterface, DetectionTestDatasetInterface
from super_gradients.training.metrics import Accuracy, Top5
from super_gradients.training import MultiGPUMode
from super_gradients.training.models.detection_models.yolo_base import YoloPostPredictionCallback
from super_gradients.training.utils.detection_utils import DetectionCollateFN
from super_gradients.training.metrics.detection_metrics import DetectionMetrics
from super_gradients.training.metrics.segmentation_metrics import PixelAccuracy, IoU


class TestWithoutTrainTest(unittest.TestCase):
    @classmethod
    def setUp(cls):
        # NAMES FOR THE EXPERIMENTS TO LATER DELETE
        cls.folder_names = ['test_classification_model', 'test_detection_model', 'test_segmentation_model']

    @classmethod
    def tearDownClass(cls) -> None:
        # ERASE ALL THE FOLDERS THAT WERE CREATED DURING THIS TEST
        for folder in cls.folder_names:
            if os.path.isdir(os.path.join('checkpoints', folder)):
                shutil.rmtree(os.path.join('checkpoints', folder))

    @staticmethod
    def get_classification_trainer(name=''):
        trainer = Trainer(name, model_checkpoints_location='local')
        dataset_params = {"batch_size": 4}
        dataset = ClassificationTestDatasetInterface(dataset_params=dataset_params)
        trainer.connect_dataset_interface(dataset)
        trainer.build_model("resnet18_cifar")
        return trainer

    @staticmethod
    def get_detection_trainer(name=''):
        dataset_params = {"batch_size": 4,
                          "test_batch_size": 4,
                          "dataset_dir": "/data/coco/",
                          "s3_link": None,
                          "image_size": 320,
                          "test_collate_fn": DetectionCollateFN(),
                          "train_collate_fn": DetectionCollateFN(),
                          }

        trainer = Trainer(name, model_checkpoints_location='local',
                          multi_gpu=MultiGPUMode.OFF,
                          post_prediction_callback=YoloPostPredictionCallback())
        dataset_interface = DetectionTestDatasetInterface(dataset_params=dataset_params)
        trainer.connect_dataset_interface(dataset_interface, data_loader_num_workers=4)
        trainer.build_model('yolox_s')
        return trainer

    @staticmethod
    def get_segmentation_trainer(name=''):
        shelfnet_lw_arch_params = {"num_classes": 5, "load_checkpoint": False}
        trainer = Trainer(name, model_checkpoints_location='local', multi_gpu=False)

        dataset_interface = SegmentationTestDatasetInterface()
        trainer.connect_dataset_interface(dataset_interface, data_loader_num_workers=8)
        trainer.build_model('shelfnet34_lw', arch_params=shelfnet_lw_arch_params)
        return trainer

    def test_test_without_train(self):
        trainer = self.get_classification_trainer(self.folder_names[0])
        assert isinstance(trainer.test(silent_mode=True, test_metrics_list=[Accuracy(), Top5()]), tuple)

        trainer = self.get_detection_trainer(self.folder_names[1])

        test_metrics = [DetectionMetrics(post_prediction_callback=trainer.post_prediction_callback, num_cls=5)]

        assert isinstance(trainer.test(silent_mode=True, test_metrics_list=test_metrics), tuple)

        trainer = self.get_segmentation_trainer(self.folder_names[2])
        assert isinstance(trainer.test(silent_mode=True, test_metrics_list=[IoU(21), PixelAccuracy()]), tuple)

    def test_test_on_valid_loader_without_train(self):
        trainer = self.get_classification_trainer(self.folder_names[0])
        assert isinstance(trainer.test(test_loader=trainer.valid_loader, silent_mode=True, test_metrics_list=[Accuracy(), Top5()]), tuple)

        trainer = self.get_detection_trainer(self.folder_names[1])

        test_metrics = [DetectionMetrics(post_prediction_callback=trainer.post_prediction_callback, num_cls=5)]

        assert isinstance(trainer.test(test_loader=trainer.valid_loader, silent_mode=True, test_metrics_list=test_metrics), tuple)

        model = self.get_segmentation_trainer(self.folder_names[2])
        assert isinstance(model.test(test_loader=model.valid_loader, silent_mode=True, test_metrics_list=[IoU(21), PixelAccuracy()]), tuple)


if __name__ == '__main__':
    unittest.main()